计算瓶颈分析：区分计算密集型与访存密集型任务

做吞吐优化，第一步不是调代码，而是搞清楚你的程序到底卡在哪里。我见过太多人上来就调算子、改并行策略，结果发现瓶颈根本不在计算上——白忙活一场。

今天咱们聊聊怎么区分两类任务：计算密集型（Compute-bound）和访存密集型（Memory-bound）。说白了，前者是「算不过来」，后者是「数据喂不饱计算单元」。

1. 两类瓶颈的本质区别

先看一个直观的例子。假设你有一个矩阵乘法：

# 计算密集型：大矩阵乘法
C = A @ B  # A: 4096x4096, B: 4096x4096

# 访存密集型：逐元素操作
D = A + B  # A: 4096x4096, B: 4096x4096

矩阵乘法需要做大量浮点运算，计算时间远大于数据搬运时间。而逐元素加法，计算量极小，大部分时间花在从显存读数据、写回结果上。

我个人习惯用两个指标来区分：

计算密度（Arithmetic Intensity） = 总浮点运算次数 / 总数据访存量（FLOPs/Byte）
硬件峰值：计算峰值（TFLOPS）和显存带宽（GB/s）

当计算密度高于某个阈值时，程序受计算能力限制；低于阈值时，受显存带宽限制。这个阈值就是Roofline模型中的「ridge point」。

2. Roofline模型：一张图看懂瓶颈

Roofline模型其实不复杂。它把程序性能画在一张图上：

横轴：计算密度（FLOPs/Byte）
纵轴：实际性能（GFLOPS）
天花板：计算峰值（水平线）和显存带宽（斜线）

我画了一张示意图，帮你直观理解：

图中红色斜线区域是「访存受限」——你被显存带宽卡住了。蓝色水平线区域是「计算受限」——GPU的计算单元在全力跑。两条线的交点就是Ridge Point。

核心判断逻辑：

程序的计算密度 > Ridge Point → Compute-bound
程序的计算密度 < Ridge Point → Memory-bound

3. 实战：如何计算计算密度

拿一个实际例子算算。假设你在做Layer Normalization：

def layer_norm(x, gamma, beta):
    # x: [B, T, D] = [4, 1024, 768]
    mean = x.mean(dim=-1, keepdim=True)      # 读x，写mean
    var = ((x - mean) ** 2).mean(dim=-1)      # 读x和mean，写var
    out = (x - mean) / torch.sqrt(var + eps)  # 读x,mean,var，写out
    out = out * gamma + beta                  # 读out,gamma,beta，写out
    return out

算一下：

总FLOPs：约 5 × B×T×D = 5 × 4×1024×768 ≈ 15.7M FLOPs
总访存量：读x(12MB) + 读mean(0.016MB) + 读var(0.016MB) + 写out(12MB) + 读gamma/beta(0.006MB) ≈ 24MB
计算密度 = 15.7M / 24MB ≈ 0.65 FLOPs/Byte

嗯，这个值很低。典型的Memory-bound任务。我在项目中遇到过类似情况，LayerNorm在A100上只能跑到理论带宽的30%左右，因为计算太简单，数据搬运成了瓶颈。

4. 常见任务的分类

任务类型	典型算子	瓶颈类型	优化方向
矩阵乘法	MatMul, Conv	Compute-bound	使用Tensor Core，增大tile size
逐元素操作	ReLU, Add, Mul	Memory-bound	算子融合，减少访存次数
规约操作	Softmax, LayerNorm	Memory-bound	使用warp级规约，共享内存
注意力机制	FlashAttention	两者混合	分块计算，减少HBM访问

我的经验：Attention机制很有意思。QK^T矩阵乘法是Compute-bound，但Softmax是Memory-bound。FlashAttention之所以快，就是因为它把这两步融合了，让计算密度大幅提升。

5. 避坑指南

我曾经犯过一个低级错误。当时优化一个BERT推理模型，发现整体性能上不去。我盯着矩阵乘法调了半天，各种tiling、double buffering都试了，提升不到5%。

后来用Roofline模型一分析，发现瓶颈根本不在矩阵乘法——Embedding层和LayerNorm占了40%的时间，而这些全是Memory-bound操作。我花了两周时间做算子融合，把Embedding + LayerNorm + Dropout合并成一个kernel，性能直接翻倍。

注意：不要只看单个算子的Roofline。整个pipeline的瓶颈可能藏在你不注意的地方。我建议你每次优化前，先用Nsight Systems或PyTorch Profiler跑一遍，看看每个算子的耗时占比。

6. 如何用Roofline指导优化

知道了瓶颈类型，优化方向就很清晰了：

Compute-bound任务：想办法提高计算效率。用Tensor Core、增大batch size、使用混合精度训练。说白了就是让计算单元别闲着。
Memory-bound任务：减少访存次数。算子融合、使用共享内存、优化数据布局。你想想看，既然计算单元在等数据，那就少搬点数据。

举个例子，如果你在做GELU激活函数：

# 原始：两次访存
x = torch.nn.functional.gelu(x)  # 读x，写x

# 融合后：一次访存
# 把GELU和前面的MatMul融合成一个kernel
# 读一次x，计算完直接写结果

融合后访存量减半，对于Memory-bound的GELU来说，性能提升非常明显。

总结一下：

用Roofline模型定位瓶颈类型
Compute-bound → 提高计算效率
Memory-bound → 减少访存次数
别凭感觉优化，先分析再动手

好了，这一章就到这里。记住一句话：优化之前先问自己——我的程序到底卡在哪里？

📚 课程目录

计算瓶颈分析：区分计算密集型与访存密集型任务

1. 两类瓶颈的本质区别

2. Roofline模型：一张图看懂瓶颈

3. 实战：如何计算计算密度

4. 常见任务的分类

5. 避坑指南

6. 如何用Roofline指导优化