3. AI工作负载特征分析
做AI芯片优化这么多年，我有个习惯——拿到一个新架构，先不看它的算力有多猛，而是先看它要跑什么负载。为什么？因为芯片设计说白了就是为应用服务的。你连CNN、Transformer这些模型的数据流特征都没摸清楚，谈什么缓存优化？
这一章，咱们就来拆解三种最常见的AI工作负载：CNN、Transformer中的自注意力机制，还有矩阵乘法。我会结合我踩过的坑，把它们的访存特征和数据复用性讲透。
3.1 卷积神经网络（CNN）的数据流模式
CNN大家都熟，但它的数据流模式其实挺有讲究的。我当年第一次做CNN加速器时，就因为在数据复用上没想清楚，导致缓存命中率惨不忍睹。
CNN的核心操作是卷积。一个卷积层，输入是特征图（IFMap），输出是特征图（OFMap），中间有卷积核（Weight）。数据流模式可以归纳为三种：
- 输入驻留（Input Stationary）：把输入特征图尽量留在缓存里，反复使用。适合输入通道数大的情况。
- 权重驻留（Weight Stationary）：把卷积核权重留在缓存里，适合权重复用度高的场景。我习惯在模型推理时用这种。
- 输出驻留（Output Stationary）：把部分和累加留在寄存器或缓存里，减少写回次数。训练时我常用这个。
你想想看，一个3x3的卷积核，在输入特征图上滑动时，相邻窗口之间有大量重叠。这个重叠区域的数据会被反复读取。如果缓存设计不合理，这些数据就会反复从DRAM里搬，功耗和延迟都上去了。
关键洞察：CNN的数据局部性主要体现在空间局部性上。相邻像素、相邻通道、相邻输出位置，都有很强的数据复用关系。
举个例子，一个典型的卷积层：输入是224x224x3，卷积核是3x3x3x64，输出是224x224x64。你算算看，每个输入像素会被多少个卷积窗口用到？答案是9次（3x3核）。如果缓存能把这9次复用都覆盖掉，带宽需求就能降一个数量级。
我的经验：在实际项目中，我建议优先考虑权重驻留模式。因为权重的复用度最高（所有输出位置共享同一组权重），而且权重通常可以预取到缓存里。我曾经在一个CNN加速器项目中，通过把权重缓存从16KB扩大到32KB，性能提升了40%。
3.2 Transformer中的自注意力机制访存特征
Transformer火了之后，很多做AI芯片的同行都跑来问我：这玩意儿跟CNN有啥不一样？我说，区别大了去了。
自注意力机制的核心是计算Q、K、V三个矩阵，然后做softmax。它的访存特征跟CNN完全不同：
- 全局依赖：每个输出位置都要跟所有输入位置做交互。这意味着数据局部性很差——你没法像CNN那样只缓存一个局部窗口。
- 中间结果多：Q、K、V、注意力分数、加权和……每一步都会产生大量中间数据。这些数据写回DRAM再读出来，代价极高。
- 序列长度敏感：序列长度一长，注意力矩阵就是O(n²)的复杂度。我见过一个项目，序列长度从512涨到2048，访存带宽直接爆了。
为什么会这样？说白了，自注意力机制的计算模式是「所有对所有的」。每个都要跟其他所有算相似度。这跟CNN那种「局部滑动窗口」完全不是一个路子。
避坑指南：我曾经在一个Transformer推理芯片的设计中，忽略了中间结果的缓存策略。结果发现，注意力分数矩阵的写回和读取占用了超过60%的带宽。后来我们不得不专门加了一个SRAM来缓存中间结果，才把性能拉回来。
我个人的建议是，对于Transformer的访存优化，重点放在两点：一是减少中间结果的写回次数（比如用在线softmax），二是利用稀疏性（很多注意力分数其实接近0，可以跳过）。
3.3 矩阵乘法的数据复用性分析
矩阵乘法是AI计算的基石。CNN里的卷积可以转化成矩阵乘法，Transformer里的自注意力也是矩阵乘法。所以，搞懂矩阵乘法的数据复用性，就等于搞懂了一半的AI芯片优化。
一个矩阵乘法 C = A × B，其中A是M×K，B是K×N，C是M×N。数据复用性可以从三个维度看：
| 复用维度 | 复用次数 | 说明 |
|---|---|---|
| A的行复用 | N次 | A的一行会被B的所有列使用 |
| B的列复用 | M次 | B的一列会被A的所有行使用 |
| 部分和复用 | K次 | 累加过程中，部分和会被反复更新 |
你想想看，如果M、N、K都很大（比如4096），那么A的每一行要被复用4096次，B的每一列要被复用4096次。这复用性其实很高！但问题在于，缓存能不能装得下？
我习惯用分块（tiling）来解决这个问题。把大矩阵切成小块，让每个小块能塞进缓存里。比如，把A切成M_tile×K_tile，B切成K_tile×N_tile，这样每个小块的大小就是M_tile×K_tile + K_tile×N_tile。只要这个和小于缓存大小，就能实现数据复用。
核心公式：缓存命中率 ≈ 1 - (分块边界数据量 / 总数据量)。分块越小，命中率越高，但分块太多又会引入循环开销。这是个trade-off。
我记得有一次优化一个矩阵乘法算子，M=N=K=2048。一开始没做分块，缓存命中率只有30%。后来做了16x16的分块，命中率直接飙到95%以上。嗯，这就是数据复用性分析的价值。
我的建议：在做矩阵乘法优化时，先算清楚每个维度的复用次数，然后根据缓存大小选择合适的分块策略。不要一上来就调代码，先想清楚数据流。
3.4 三种负载的对比总结
说了这么多，咱们来做个对比。我把三种负载的核心特征整理成了一张图，方便你理解。
这张图把三种负载的核心特征都列出来了。你可以看到，CNN和矩阵乘法在数据复用性上其实有共通之处——都适合用分块策略。而Transformer则更依赖中间结果的缓存和稀疏优化。
好了，这一章的内容就到这里。记住一句话：没有通用的缓存优化方案，只有针对特定负载的定制策略。搞清楚了数据流和复用性，你离一个优秀的AI芯片架构师就不远了。