实在之路
黑格尔✅Cuda Reduce优化笔记
type
status
date
slug
summary
tags
category
icon
password
AI summary
对cuda_samples中reduction优化case的个人学习笔记，学习cuda以及nv的架构。
参考:
https://developer.download.nvidia.com/assets/cuda/files/reduction.pdf
https://zhuanlan.zhihu.com/p/426978026
下面贴出来的代码均出自：
https://github.com/NVIDIA/cuda-samples/tree/master/Samples/2_Concepts_and_Techniques/reduction
Tree
gpu上，reduce可用树形递归做
上图是一个thread block的任务。
但是需要在thread block之间通信， nv没有这种手段。
为什么没有：
- 硬件限制，在处理器数量很多的条件下，硬件实现这个功能也很难
Multiple Kernels
解决方法： 拆成多个kernel。
- Launch kernel的作用相当于是一次global sync
上面两次launch 的kernel是一模一样的。
优化目标
- 选择合适的metric
- GFLOP/s
- Bandwidth
- Reduction的计算密度很低
- 1 flop per element overloaded
- 因此我们关注Bandwidth
base line reduce0
- 取模运算非常慢， 可以将
tid % (2*s)替换为tid &(2*s -1), 因为这里s是2的幂次。
wrap diverge- 每个wrap，第一轮只有 1/2 的thread 是 active的
- 第二轮 1/4 ，... 1/8, ....1/16, .... 1/32
- 第六轮：一半的wrap的全是inactive的， 1/2 的wrap只有一个 active thread
- ...
单看一个wrap, 这里的diverge会导致wrap的执行时间增加吗？
这里只有if 分支，而没有分支，逻辑上，diverge和所有thread都走if 分支需要的时间不是一样的吗？
Fix wrap diverge reduce1
让同一个wrap内的thread，都尽量走一个分支，上图中，第一轮时，工作的tid为：
0, 2, 4, ...
我们挤压一下，让工作的tid变为:
0, 1, 2..
这样在第一轮，前面一半的wrap都会进入if分支，而剩下的一半不会。
for循环改写为：
- Bank conflict，考虑t0 t16, load的第一个数据，index分别为0,32, 位于同一个bank。
- 从总体来看，到了后续几轮后，wrap内依然会存在inactive thread
Fix bankconflict reduce2
让不同线程的同一条load指令，load的数据连续。
for循环改写为
Fix idle thread reduce3
第一轮reduce中，有一半的thread / wrap只进行load，不进行compute, 浪费了硬件资源。
能不能让这些wrap提前return？ 结束运行，释放硬件资源？
thread数量减半，每个thread，进来时load两个数据，并做一次加
这里把第一轮reduce提了出来。 可以理解为：第一轮待计算数据有很多，不应该有thread空闲，所以显示地把这个动作个每个thread加上了。
按照这个逻辑，如果数据规模是在编译期已知的一个值，是不是可以用模板元编程的技术，在编译期展开，让每个thread在进入for loop前，完成多轮的reduce？其实就是reduce6的优化。
Use Wrap Shuffle reduce4
最后一个wrap时，每个thread读的数据都来自同一个wrap内的其它thread， 现在一个wrap内可以通过一个指令直接访问其它thread的寄存器，无需经过smem。
改动： for loop不进行最后一个loop，也就是最后一个wrap内的计算
最后一个wrap的计算，通过wrap内通信完成
Unroll for-loop reduce5
- reduce时计算密度很低
- bottleneck很可能是 instruction overhead， 即：
- 不是load，不是store，不是reduce 计算的那些instruction
- 也就是： 地址计算, fo循环的overhead
方法： 展开for循环，没什么好说的
上面blockSize是编译期常量，相关if都会在编译期确定。
能不能通过模板元编程，编译期递归展开，让kernel更具泛化性。
Add multiple elements reduce6
分析下之前kernel的复杂度,我们只分析一个block内部的复杂度，
假设一个block处理的数据量是N
观察上面的kernel， kernel可以抽象成3部分：
- load部分， 复杂度为： 1
- 树形递归计算部分, 复杂度：logN
- Store 部分，1
因此单个kernel的复杂度为O(logN),
一个wrap有32个thread，并且一个sm有4个processor，再加上wrap scheduler的hide latency的效果，一个block内的thread之间有一定的并行能力。
但这依然是有限的，我们不能认为所有thread是完全并行的。因此整体的时间复杂度是 高于 logN的。具体的复杂度，取决于很多因素: hardware spec, 一个wrap对资源的需求(比如register等）...
NV的slide中，基于G80架构。认为复杂度是 N*logN，并以此来分析。
改进办法，修改算法流程如下， 假设我们每个block只用P个thread， 每个thread就需要load P/N的数据。
- load N/P的数据， 复杂度为： N/P
- 每个thread对 load进来的 N/P个数据，做累加,得到一个结果，放在shared mem上。 复杂度: N/P
- 树形递归计算部分, 复杂度：logP。 （只剩下P个数据了）
- Store 部分，1
总复杂度： P*(N/P + N/P + logP) = N + P*logP
我认为这个改进本质上是reduce3 的plus版本。
实测这个优化在4090上，把throughput从400G提升到了800G。
Seq Wrap Wrap reduce7
基于上面的reduce6， 当P个thread做完自己的那部分N/P的累加后，每个thread都有一个数据，此时可以先在每个wrap做reduce, 每个wrap得到一个数据，放到smem上。 由于P 不会超过1024 (cuda 限制), 且 wrap_size =32， 所以这一步剩下的数据不会超过32，可以用一个wrap完成。
流程如下：
- load N/P的数据， 复杂度为： N/P
- 每个wrap内部，使用wrap shuffle指令，高效地做 warpReduceSum，并把结果放到smem[tid % wrap_size]上。 复杂度： log32
- 最后一个wrap再做一次 warpReduceSum. 复杂度：log32
- Store 部分
使用这种方法thread，只需要一次wrap间的同步，即第一次warpReduceSum后。除了边界情况，也不会有wrap diverge。
4090上的实测性能情况
4090， compute capability 8.9
上面的reduce0-reduce5，在4090上的差距都不明显，throughput都在400G左右。
但reduce6直接将throughput拉到800G，达到理论上限的80%。
reduce7与reduce6差距不明显。
总结
这里的优化针对的是一个巨大的reduction, 但是实际我们的数据量不会有这么大。不过主要是通过这个case，学习优化思路以及N卡的一些硬件架构特点。
Extra: 使用Tensor core做reduce
这篇文章介绍了如何将reduce mapping到tensor core unit 上面执行。其中tensor core unit是专门做mma(matrix multiply add)的硬件模块。
即：
不同的硬件下，对MNK的值往往是有限制的，文章选取M=N=K=16作为基本shape来进行后续的op mapping.
TCU REDUCTION ALGORITHM
很好懂，就是构造一个第一行为全1的矩阵， 后续将这个矩阵叫做 P.
因为矩阵乘的定义中,本身就是k轴对应元素相乘,再累加. 所以只要将令某一个矩阵的k轴全1, 就相当于是在做reduce.
Reduction16
假设input shape为：[16, 16]，在dim 1做reduce， 则output shape为： [16, 1]
- 将input按照col major load，其实就是把需要reduce的维度，作为mma的 k 轴。
- 然后计算 ， 得到结果
Reduction256
假设input A shape为：[N*256]， 则output V shape为： [1]
- 每次load 256个数据
- 计算 ， 其中 是上一次的计算结果
- 最后通过 得到结果
Prev
TVM论文笔记
Next
RoPE 旋转位置编码
Loading...