手撕深度学习之CUDA并行规约算法（上篇）：硬核揭秘200%性能提升的GPU优化之道，从硬件特性到算法实现的完整进阶指南

本文首发于本人的微信公众号，原文链接：

https://mp.weixin.qq.com/s/nzE2NB7T2U2F46IBJmPSYQ

摘要

本文为CUDA并行规约系列文章的上篇，本系列将会介绍CUDA编程的一些基础软硬件知识，然后给出7种规约算法的实现，并从硬件的角度对它们进行分析和优化，最终给出一个开箱即用的模板代码。

本文主要介绍了CUDA编程的基础软硬件知识，并给出了2种规约算法的实现和分析。

写在前面

最近在写DL Systems的Homework 3，遇到了要使用cuda来实现reduction（例如sum，max等）的任务。几年前写cuda时我也曾被这个问题难住过，于是便正好借此机会深入研究一下其中的门道。

要实现一个简单的reduction算法并不难，只需要把数据分块扔给执行线程，然后每个线程里用for循环做reduction即可，由于有GPU并行处理的助力，这一实现肯定是会比CPU上更快的。

但是如果想要更进一步地利用上GPU的能力，就需要花一些功夫来设计这个算法了。这其中涉及到了GPU的若干硬件特性，包括SIMT（Single Instruction, Multiple Threads）架构，Wrap（线程束），Bank（存储体）等。NVIDIA甚至专门出了个PPT来讲解如何在GPU上高效地完成reduction，本系列也是以这个PPT为蓝本展开的。

由于篇幅原因，本系列预计会分为上下两篇，上篇主要介绍一些前置知识，基本思想和两个比较初级的reduction算法；下篇会接着介绍另外五种升级版的reduction算法，并给出一个使用C++模板实现的，可以开箱即用的reduction算法实现。

希望本系列文章能够给大家带来帮助。在GPU并行计算方面我也是刚入门的小白，如果文章中有疏漏之处，还请大佬们批评指正。

参考资料

NVIDIA的并行计算PPT：
https://developer.download.nvidia.cn/assets/cuda/files/reduction.pdf

7 Step Optimization of Parallel Reduction with CUDA：
https://medium.com/@rimikadhara/7-step-optimization-of-parallel-reduction-with-cuda-33a3b2feafd8

关于GPU你需要知道的小知识

这一章节主要介绍一些GPU硬件层面的前置知识。CUDA算子开发是一个需要尽可能榨干硬件全部性能的工作，所以这就要求开发者对硬件有着一些基本的了解。在具备了硬件相关的基础知识后，我们就能轻松的回答如下问题：

• 为什么仅仅是去掉一个if-else，算子性能就提升了一倍多？
• 为什么只是改了一下元素访问顺序，计算速度就快了那么多？
• ......

CUDA编程基础

这里再简单回顾一下CUDA编程里的一些基础概念。对这部分比较熟悉的朋友可以直接跳到下一小节。

内核函数

我们开发者会使用C++编写内核函数，然后通过在外部调用内核函数的方式来启动这个内核函数。

在启动内核后，会有若干个处理器并行地执行相同的 内核函数代码。开发者可以在内核函数里通过blockIdx和threadIdx来知道当前执行的线程的编号，从而对内存中的指定区域进行操作。

例如，如下是一个经典的内核函数的例子，这个例子演示了向量的加法操作：

__global__ void VectorAddKernel(const float *a, const float *b, float *out, size_t array_size) {
  const size_t idx = threadIdx.x;
  if (idx >= array_size) {
    return;
  }

  out[idx] = a[idx] + b[idx];
}

由此也可以窥见CUDA编程的风格：即使用相同的代码来操作数据中不同的区域，并且内核函数不返回结果，而是把结果写入到某个地址中。

线程组织架构

紧接着上面的向量加法的例子，那么blockIdx和threadIdx是怎么来的呢？答案是由调用这个内核函数的主机函数指定的。下面是调用向量加法内核函数的主机函数的例子：

// 假设这里的a, b, out的地址都是CUDA地址
void VectorAdd(const float *a, const float *b, float *out, size_t array_size) {
  dim3 grid(1, 1, 1);
  dim3 block(array_size, 1, 1);
  VectorAddKernel<<<grid, block>>>(a, b, out, array_size);
}

主机函数通过传入的grid和block来指定会有多少个线程参与运算。具体而言，线程组织架构分为3级，分别是Grid, Block, Thread，其中Thread是真正做事的，若干个Thread组成一个Block，若干个Block组成Grid，目前可以认为全局只有一个Grid。

为了编程方便，我们把Grid和Block都分成了xyz三个维度，如下图所示：

上面代码里传入的实际上是每个维度里Block和Thread的数量。在上面的例子中，一共有\(1 \times 1 \times 1 = 1\)个Block，每个Block里有\(array\_size \times 1 \times 1 = array\_size\)个线程。

需要注意的是，虽然这里使用xyz分成了三个维度，但是在硬件层面上，所有线程都是等价的，所以这里的分维度仅仅是为了在逻辑上表示起来更方便，实际在执行时维度并不会产生影响。

不互通的内存

再回到上面的主机函数中，第一行有个注释："假设这里的a, b, out的地址都是CUDA地址"。这就涉及到了CUDA编程中的又一个注意事项：虽然内核函数和主机函数都是写在同一个文件里的，但是实际执行时，内核函数会运行在GPU上，而主机函数还是在CPU上，所以两者的地址是不互通的 。也就是说，主机函数上的一个地址，对应的数据是在内存条 中的，如果不加以转换，就把这个地址传递给内核函数，那么内核函数就会尝试在显卡内存中对应位置去取数据，就会造成严重的错误。

具体的转换方法是使用cudaMalloc申请内存，然后再把主机内存中的数据copy到显卡内存中。在数据量大的时候，这也会是一个不可忽视的瓶颈。因此，pytorch等框架会让数据尽量一直放在CUDA中，以避免数据转移的开销。

SIMT模型与线程束

GPU在实际进行调度时是以一个线程束（Wrap）为单位进行调度的，一个Wrap包含32个线程，其执行过程遵循SIMT（Single Instruction, Multiple Threads）架构的设计。

具体而言，GPU在取指令后会把该指令广播给Wrap中所有的32个线程，这32个线程会同时执行该指令，然后在下一个时钟周期到来后，再重复相同的过程。所以Wrap中的线程可以认为一定是同步的，因为同一时刻每个线程都在执行相同的指令，这是硬件上保证的。

那么这时候问题就来了，如果遇到if语句，导致一部分线程和另一部分线程要执行的指令不同会发生什么呢？

答案是这时候就会发生所谓的分支发散（Wrap Divergence），在发生Wrap Divergence时，这些线程的执行会由并行转成串行，GPU会首先执行第一个分支的代码，那些满足第一个分支条件的线程会被激活，其余线程会被冻结，执行完第一个分支后，再执行第二个分支，同样冻结一部分线程，激活一部分线程，直到分支结束。

Wrap Divergence会导致严重的性能损失，因此在做CUDA开发时一定要尽可能的保证一个Wrap里所有线程都执行相同的代码。（这一点在后面具体的reduction算法部分也会涉及到）

线程同步

CUDA支持一个线程块（Block）内 进行线程同步。具体而言，只需要调用__syncthreads，调用之后线程就会阻塞，直到线程块内所有线程都执行到这一点。

共享内存与Bank Conflict

我们在主机函数中使用cudaMalloc申请的内存是位于GPU板载DRAM上的全局内存，全局内存的容量大，但是访问慢。那么相对应的，也会有一块容量小，但是访问快速的内存，它就是共享内存（Shared Memory）。共享内存位于GPU的流处理器的SRAM上，其作用域仅限于Block内。

为了提升访问速度，GPU把共享内存分为了32个Bank，这里的Bank可以理解为口岸的意思，也就是把对内存的访问分摊到了32个口岸来完成，这些口岸是可以并行工作的，访问一个地址只能去该地址对应的口岸访问。

具体而言，访问地址0需要去第0个Bank，访问地址1需要去第1个Bank，访问地址31需要去第31个Bank，访问地址32则又回到了起点，需要去第0个Bank，以此类推。

理想情况下，一个Wrap里的线程会按顺序访问0，1，2...，31号地址，那么这32个访问在一个时钟周期内就能完成。但是如果一个Wrap里有2个线程访问了地址0，或者一个访问了地址0，另一个访问了地址32，就会在Bank 0上发生Bank Conflict。

在发生Bank Conflict时，由于GPU无法在一个时钟周期内同时处理这两个互相冲突的访问请求，所以就只能把这两个请求串行化，导致整个Wrap的执行节奏被拖慢。最坏情况下，如果32个线程都访问同一个Bank，那么这个Wrap就需要花费原来32倍的时间来执行这个访存指令。

Bank Conflict同样会导致严重的性能损失，这也是我们在开发时需要注意的一点，这一点在后面具体的reduction算法实现中也会提及到。

问题定义

这里先严格定义一下我们要解决的问题：我们有一个矩阵\(X \in \mathbb{R}^{m \times n}\)，一个规约函数\(Reduce(x_1, x_2, \ldots, x_n)\)。我们需要求一个向量\(O \in \mathbb{R}^{m}\)，满足\(O_i = Reduce(x_{i1}, x_{i2}, \ldots, x_{in})\)。

特别的，我们还要求Reduce函数满足如下性质：\(Reduce(x_1, x_2, x_3) = Reduce(Reduce(x_1, x_2), x_3)\)。

举个例子：这个Reduce函数可以是求和函数Sum，为了便于理解，后面我们就直接以Sum为例展开了。

这个问题换一种描述就是：输入\(m \times n\)的矩阵，然后对每一行进行求和，输出求和之后的结果。

需要注意的是，在具体实现中，我们会把矩阵X按行优先的方式展开成一个\(m \times n\)的一维向量。因为在C++中实现多维数组比较麻烦，一般在底层都是使用一维数组+维度信息的方式来表示一个多维数组的。

算法基本思想

我们可以认为是有m个batch，每个batch里的n个元素做Sum操作。因此可以开m个Block，每个Block里有n个Thread。由于每个Block里的操作都相同，所以我们就把视角聚焦到某个具体的Block。

算法大致的流程为：首先n个线程会把n个数据复制到共享内存中，每个线程都对应和负责共享内存中的一个位置 ，然后再做类似于从下至上的归并排序的操作，即\(T_0\)把\(A_0\)和\(A_1\)的数值相加，然后存储到\(A_0\)中；\(T_2\)把\(A_2\)和\(A_3\)的值相加，存到\(A_2\)中，以此类推（注：这里\(T_0\)表示编号为0的线程，\(A_0\)表示共享内存中编号为0的元素）。

至于\(T_1\), \(T_3\)这些线程，它们不会做任何事。

通过块内的线程同步机制（__syncthreads）确保所有线程都执行完成后，就进入下一轮归并，此时就是\(T_0\)把\(A_0\)和\(A_2\)的数值相加，存到\(T_0\)；\(T_4\)把\(A_4\)和\(A_6\)相加，存到\(A_4\)......。而其余线程，如\(T_1\), \(T_2\), \(T_3\)，则什么也不做。

如此循环下去，直到最后步长为n，此时\(A_0\)存储的就是全部n个数的和，此时由\(T_0\)把\(A_0\)写入到输出向量的对应位置即可完成任务。

算法具体实现（上篇）

为了便于讨论，这里就取n=256，即n是2的整数次幂。这里其实有一个open question：即在实际实现时是否需要考虑到n不是2的整数次幂的情况，做边界条件的处理？这里个人认为不应该考虑边界条件，应该由上层来做padding来保证n一定是2的整数次幂。因为如果需要考虑边界条件，就不可避免的会引入很多if-else，这会增大触发Wrap Divergence的概率，比如偶尔会需要对奇数值做一些修正等等。

具体的padding方式需要看Reduce函数是什么，如果是Sum，可以直接填0，如果是Max，可以考虑填-inf等。

总之，个人的观点是：CUDA内核函数不应该过多的处理边界条件，边界条件的处理应该尽可能由上层来完成。

V0版本：Interleaved Addressing

这一版本的算法就是上面算法基本思想部分的实现，具体的内核函数代码如下图所示：

具体的执行过程可以通过下面这张图看出：

从\(T_0\)的视角来看，第一个循环里，计算\(A_0 = A_0 + A_1\)，第二个循环里，计算\(A_0 = A_0 + A_2\)，由于此时在上一轮中，\(T_2\)已经完成了\(A_2 = A_2 + A_3\)，所以这一轮循环结束时，\(A_0\)就已经完成了前4个元素的加和。

在能够确保n为2的整数次幂的情况下，整个算法实现就会非常简单，只需要一路把stride *= 2即可，无需考虑任何边界条件。

V0版本的不足之处

V0版本主要有2个问题：

1. 取模操作非常耗时，因为除法操作在硬件底层实现就是会慢一些
2. Wrap Divergence：我们可以注意到，在第一轮循环中，\(T_0\)需要执行加和操作，所以是走的if里的true分支，而\(T_1\)不需要任何操作，所以走的是false分支，由于两者位于同一个Wrap里，所以此时已经造成了Wrap Divergence，这会严重影响执行性能。

V1版本：Interleaved Addressing（改进版）

有一个可以同时修复掉取模和Wrap Divergence的方法，具体修复方法如下：

这一算法的主要思想是：从后往前让线程停止工作 。具体而言，以n=256为例，在第一次循环时，只有\(T_0\)到\(T_{127}\)在工作，而\(T_{128}\)到\(T_{255}\)会走到if的false分支，什么也不做；在第二次循环时，\(T_{64}\)到\(T_{127}\)又会走到false分支，以此类推。更具体的可以看下面这张图

这一算法保证了在工作的线程和在休息的线程始终是扎堆的，从而避免了一开始就Wrap Divergence的局面。

NVIDIA也在PPT里面展示了这一优化的效果，如下图所示，还是相当可观的：

V1版本的不足之处

首先，这个方法肯定是没有彻底解决Wrap Divergence的，毕竟到最后就只有\(T_0\)在工作了。当然，这也是一个优化方向，这个后面会提及到。

其次，这一方法还有可能导致Bank Conflict。观察\(T_0\)和\(T_1\)在每次循环中的访问内存地址，可以发现两者差距越来越大，并最终会在某一时刻，\(T_1\)要计算\(A_{32} = A_{32} + \ldots\)，而此时\(T_0\)也要访问\(A_0\)，这就导致了Bank Conflict，因为两者都同时想要访问Bank 0。

并且更糟糕的是，此时\(T_0\)的Wrap里的所有线程都会访问Bank 0，这是Bank Conflict里最糟的情况，会导致Wrap内的并行完全变成串行。

至于这一问题的解决办法，还请听下回分解啦。