
AI infra:内存模型与规约优化
分析了多个线程同时累加同一变量时产生的数据竞争问题,并引出原子操作的作用与性能瓶颈;随后通过 warp 级规约、block 级 shared memory 规约、__shfl_down_sync() 线程束洗牌、Two-Pass Reduce 和 Cooperative Groups 等方法,逐步优化跨线程、跨 block 的求和过程。最后结合 GPU 内存层次结构与 Bank Conflict 问题,说明共享内存访存模式对规约性能的影响,帮助理解 CUDA 中并行规约从正确性到高性能实现的完整优化路径。
从并行加法到并行累加
之前学习的并行加法是一种尴尬并行(Embarrassingly Parallel),任务之间完全独立且没有任何数据依赖关系,也不需要线程之间的通信与同步。
但是如果换成另外一种加法,即对单个变量进行累加求和呢?
设定 input 数组中的数字全部为 1,累加为一个数,运行得到的结果如下:
结果是错误的。因为这是并行,多个线程同时进行写入操作,计算实际上会被分为三步,首先从全局内存中把 *result 读取出来,将 input[idx] 加上去,再将新的值写回 *result,多个线程会同时读到一个旧的值, 然后各自相加结束之后再写回去,这样以后写回的线程会将前面的线程的结果覆盖掉,导致加法丢失。
为了解决这个竞争的问题,需要确保顺序执行,同时又只能够有一个写入操作。可以通过原子操作来解决这个问题。所谓的原子操作,就是要么一次性执行完毕,要么不执行。可以通过 CUDA 提供的 atomicAdd() 来执行原子累加操作:
运行得到的结果如下:
可以看到答案对了,但是运行时间却多了许多。通过查看 ncu 的对比分析,可以看到计算与内存吞吐/利用率都非常低,很多线程都在闲置:
原子操作导致线程之间完全串行,看似在“并行”,实际上本质上还是串行,这就导致了运行时间变长。
并行累加——分治规约
为了解决原子操作导致线程之间的实际串行行为,可以使用分治规约的方法来解决。分治规约就类似于排序算法中的归并排序,首先先两两相加得到下一个层级结果,再对下一个层级的进行两两相加…直到得到最终的结果为止。
在 CUDA 中相应的代码示例如下:
其中 lane_id 是一个线程在所属的 warp 中的位置编号,判断其是否为 0 后再进行下一步处理,实际上就是让每一个 warp 中的 lane_id 为 0 的线程“代表”整个 warp 将应当相加的数据进行求和。而 warp 之间则通过 atomicAdd() 函数进行累加到全局的结果中。
当然这里的 32 有些许笼统,只用于案例说明,实际编写过程中可以使用 cudaDeviceProp 来获取 warpSize。在这样做之后,性能能够得到极大加强:
原来的每个线程都做原子操作导致串行化,现在每个 warp 都有一个线程在独立工作,原子操作数量从原来的线程数变成 warp 数,提升了并行度。那么是不是进一步提升并行度,规约的层级从 warp 进一步细化到 block,就能够提高呢?
运行得到的实际结果如下:
实际得到的结果甚至要比之前耗时长一些,与前面的推断不符。
通过 NCU 性能对比查看:
两者 compute 所占据的百分比都比 memory 要更高,与之前的逐元素加法不同,目前的累加实现大部分时间都在计算,并且 block 规约级别的计算利用率和活跃的 warp 数都要比 warp 级别规约的要少许多。这是因为虽然减少了原子累加的操作数,但是单个线程的计算量增加太大,导致利用率变低。这并不说明不需要 block 级别的规约,可以通过加大分治粒度的方式来进行优化。
回顾:内存墙及内存结构
学过计算机组成原理的人都知道,现有的计算机基本上都是冯·诺伊曼架构,数据必须在处理器和内存之间频繁传输,当数据量大或者计算密度低时,很容易造成内存瓶颈,并且当前的问题其实是计算的速度要远高于内存访问的速度,为了解决速度和内存访问的问题,计算机科学家们提出了阶层架构。
- 越靠近上层/处理器/计算单元的内存通常越快、容量越小、制造成本越昂贵,这类内存通常用于存放更加常用的数据和指令。
- 越靠近下层的内存越慢、越大、越便宜,这类内存通常用于存放更不常用/更大规模的数据和指令。
而 GPU 中也同样有类似的架构,相应的物理视图如下:
- 寄存器:on-chip,每个线程私有,存储线程的局部变量。
- 共享内存:on-chip,每个 SM 的所有 Block 共享,由程序员显式管理。
- L1 缓存:on-chip,每个 SM 的所有 Block 共享,由硬件自动管理。
- 只读缓存:on-chip,每个 SM 的所有 Block 共享,缓存只读数据(纹理/常量/全局)。
- L2 缓存:on-chip,全芯片所有 SM 共享,由硬件自动管理。
- 全局内存:on-GPU, off-chip,全设备的所有线程均可访问,由程序员显式管理。
- CPU 主存:off-GPU,通过 PCIe/NVLink 与 GPU 交互,由操作系统管理。
从 CUDA 的角度看,GPU 的逻辑视图如下:
- 寄存器:线程私有、速度最快、容量最小,用于保存局部临时变量。
- 共享内存:block 内线程共享的高速存储,用于线程协作和数据复用。
- 本地内存:逻辑上线程私有、物理上通常位于全局内存,用于寄存器放不下的局部数据。
- 全局内存:全设备可访问的大容量通用存储,主要存放输入输出和大规模数据。
- 常量内存:全设备可见的只读小容量存储,适合广播式读取常量参数。
- 纹理内存:带缓存的只读访问空间,适合具有空间局部性的读取模式。
回归分治规约
在了解完内存结构之后,加大分治粒度就可以通过使用共享内存来实现,数据在共享内存中对同块中的线程均为可见,可以利用共享内存做树状规约。
此处使用动态分配共享内存,__shared__ 表示该变量分配至共享内存,__syncthreads() 则是同步块内线程。代码的总体逻辑是每个线程先从全局内存读取一个元素到共享内存,而后在共享内存中进行树状规约,block 内最后只剩下最后一个和时,0 号线程会将和原子地加到全局结果中。
可以看到性能相比之前又进一步提升了。
并行累加性能继续提升
树状规约仍然存在一个问题:越往后活跃的线程越少,闲置的线程越多,利用率也就进一步下降了,原子操作现在每个 block 都会有一次,那能不能进一步减少呢?
针对线程闲置问题,可以用 线程束洗牌函数 __shfl_down_sync() 来解决,从 lane_id 和 offset 来取 sum 的值。
__shfl_down_sync() 是 CUDA 提供的一个 warp 级数据交换指令,其作用是让一个线程直接读到同一个 warp 内,编号比自己大的某个线程寄存器中的值,常见形式是 __shfl_down_sync(mask, var, delta, width),其中 mask 表示哪些 lane 可以参与这次 shuffle,写 0xffffffff 表示所有的 lane 全部参与,var 是当前线程中的寄存器变量,即想要交换的数据,delta 表示向下偏移多少个 lane 取值,如果越界就返回未定义或者原值相关行为,无须担心边界问题。因为其允许 warp 内线程直接交换寄存器中的数据,因此可以在不使用共享内存和 block 级同步的情况下完成 warp 内树状规约。通过不断缩小偏移量 offset,可逐步将一个 warp 内各线程的局部结果合并到 lane 0 中,从而减少共享内存访问和同步开销。
可以进一步抽象出使用 __shfl_down_sync() 的 warp_reduce() 函数:
从而在 warp 内使用 __shlf_down_sync,在 warp 间则使用共享内存进行规约,最终将结果加到最终输出。效果如下:
而针对每个 block 都要进行一次原子操作的问题,可以通过分级的方式来完成,分为两个核函数来分别完成规约,第一个完成 Block 内规约,第二个在第一个的基础上跨 Block 规约。
第一个核函数的结果需要存储到全局内存中,为此需要分配一个临时数组,最后的块内规约从原子操作的跨 block 规约变为了直接赋值到中间数组,第一个核函数的代码如下:
第二个核函数的逻辑和第一个核函数的逻辑基本保持一致,代码如下:
实际测试的结果如下:
可以看到性能相比于之前的结果反而下降了,这是由于分为了两个 kernel,虽然没有了原子操作导致的开销,但是多出了 kernel 发射、额外计算、访存以及隐式同步的开销。
Smem Atomic Reduce vs. Two-Pass Reduce
Smem Atomic Reduce 和 Two-Pass Reduce 都是用于跨 block 规约的常见方案。Smem Atomic Reduce 通常先在每个 block 内部使用 shared memory 完成局部规约,然后通过原子操作将每个 block 的部分结果累加到全局结果中,其优势是实现简单,只需要一个 kernel 即可完成整体规约,避免了额外的 kernel 启动和中间结果管理;但缺点是最终阶段依赖原子操作,当 block 数量较多或多个线程频繁竞争同一个全局地址时,原子操作会造成明显的串行化开销。Two-Pass Reduce 则将规约过程拆成两步:第一个 kernel 先计算每个 block 的局部结果并写入中间数组,第二个 kernel 再对这些中间结果继续规约,优点是避免了全局原子操作,性能更加稳定,尤其适合数据规模较大、block 数量较多、原子竞争严重的场景;但它需要额外的 kernel 发射,并且依赖两次 kernel 之间的隐式同步,同时还会引入中间数组的读写开销。因此,当数据量较小或对实现简洁性要求更高时,可以优先使用 Smem Atomic Reduce;而当数据量较大、原子操作竞争明显时,Two-Pass Reduce 往往更加合适。
并行累加—跨块规约优化
之前使用了两个 kernel 来进行规约,实际上也可以像同步块内/warp 内线程一样同步块间线程,前提是 CUDA 的版本为 9.0+,NVIDIA 提供了 Cooperative Group 来实现跨块的规约优化。代码如下:
代码的逻辑与之前的 two-pass 一样,其中 block.sync() 相当于 __syncthreads(),grid.sync() 为跨 block 同步,这种方式在数据量大的时候优势会更大。
相比于 Two-Pass,Grid Sync 是单个核函数,避免了额外的核函数发射开销,结合其他 CG 功能允许更多细粒度操作;但是其也较为复杂,注意事项较多,需要仔细保证前提成立,否则可能会导致死锁等行为,并且无法支持任意的 grid size。
Bank Conflict
为什么不是对称树状?
之前的树状规约是对半向前的树状规约方式,而不是以数组中点的对称树状规约方式,实际测试下来数组中点的会慢一些,引发 Bank Conflict。先把目光转移回共享内存:为了实际实现高内存带宽以支持并发访问,共享内存被划分为大小相同的内存模块(即 bank),一般为 32 个,每个 bank 为 4 bytes,连续地址的数据按照 bank 大小循环映射到不同 bank 中,同一个 warp 中的所有线程同时执行同一无冲突的指令,就能够在一个时钟周期内完成。
而Bank Conflict 指的是多个线程同时访问共享内存时,请求落到了同一个 Memory bank 上,导致本来可以并行完成的访问被拆成串行或多次执行。
如果是多个线程同时访问同一个 bank 的同一个地址时,这是就不会发生 bank conflict,而是会直接发生广播/多播。
解决 Bank Conflict
解决 Bank Conflict 的方法主要有以下两种:
一是算法/访存模式优化,这是上上策,也是最彻底、最根源性的解决方法,例如将以数组中点的对称树状规约方式优化为对半向前的树状规约方式。
二是内存填充,故意在共享数据的内存布局中插入一些无用空位,打乱原本容易映射到同一个 bank 的整齐映射关系。例如在二维共享内存矩阵转置中,假设共享内存数组定义为 __shared__ float tile[32][32],当一个 warp 的 32 个线程按列访问该数组时,例如线程 t 访问 tile[t][0],由于相邻两行之间正好间隔 32 个 float,而 shared memory 通常有 32 个 bank,这些地址在映射时就会落到同一个 bank 上,从而产生严重的 bank conflict。为了解决这个问题,可以将数组改为 __shared__ float tile[32][33],即在每一行末尾额外填充 1 个元素。这样一来,相邻两行的跨度就从 32 变成了 33,线程按列访问时对应地址对 bank 数量取模后的结果会错开分布到不同 bank,从而有效避免冲突。


