算子与数据流：典型AI算子的数据访问模式

做算子融合优化，说白了就是要搞清楚每个算子是怎么「吃」数据的。

我刚开始接触AI加速时，总觉得融合就是把几个算子拼在一起，简单粗暴。后来被现实狠狠教育了一回——有一次做MobileNet的优化，融合后性能反而下降了15%。排查了半天，才发现是没搞懂每个算子的数据访问习惯。

嗯，今天我们就来聊聊这个。

卷积（Conv）的数据访问模式

卷积算子，你想想看，它的数据访问其实挺有规律的。

一个典型的卷积操作，输入是 [N, C, H, W]，权重是 [K, C, R, S]。计算时，每个输出像素需要累加 C * R * S 次乘加。

具体来说：

• 输入特征图：同一个输入像素会被多个输出通道复用。比如3x3卷积，一个输入点最多参与9个输出点的计算。
• 权重：权重是只读的，而且访问模式很规整——按输出通道顺序遍历。
• 输出：输出是顺序写入的，每个位置只写一次。

我在项目中遇到过一个问题：当输入通道数很大（比如512）时，权重在缓存中的命中率会急剧下降。为什么？因为权重尺寸超过了L1缓存容量，每次加载新输出通道时，之前的权重就被踢出去了。

矩阵乘法（GEMM）的数据访问模式

GEMM，说白了就是 C = A * B。在AI里，全连接层、卷积的im2col实现，本质上都是GEMM。

GEMM的数据访问模式，我总结为三个字：「分块」。

一个 M x K 的矩阵A乘以 K x N 的矩阵B，得到 M x N 的矩阵C。标准实现中：

• 矩阵A：按行访问，每行K个元素。
• 矩阵B：按列访问，每列K个元素。
• 矩阵C：按行写入，每行N个元素。

这里有个大坑——矩阵B是按列访问的，而数据在内存中是按行存储的。这意味着每次访问B的一列，都会产生大量的缓存缺失。

解决这个问题的标准做法是「分块」：

// 伪代码：分块GEMM
for (int i = 0; i < M; i += BLOCK_M) {
    for (int j = 0; j < N; j += BLOCK_N) {
        for (int k = 0; k < K; k += BLOCK_K) {
            // 加载A[i:i+BLOCK_M][k:k+BLOCK_K] 到L1
            // 加载B[k:k+BLOCK_K][j:j+BLOCK_N] 到L1
            // 计算C[i:i+BLOCK_M][j:j+BLOCK_N] += A * B
        }
    }
}

分块大小怎么选？我个人习惯是让 BLOCK_M * BLOCK_K + BLOCK_K * BLOCK_N 刚好小于L1缓存大小。这样A和B的分块都能留在缓存里，反复使用。

逐元素算子（Element-wise）的数据访问模式

逐元素算子，比如ReLU、Sigmoid、加法、乘法，这些家伙的数据访问模式最简单——一对一。

每个输出元素只依赖一个（或两个）输入元素，没有复用。说白了就是「读一个，算一个，写一个」。

这种模式的好处是：

• 数据流非常规整，没有乱序访问
• 计算密度低，瓶颈通常在内存带宽
• 非常适合和前面的算子融合——因为不需要额外的数据复用

你想想看，如果一个卷积后面跟着一个ReLU，正常流程是：卷积写结果到DRAM，ReLU再从DRAM读出来。这多浪费啊！

三种算子的对比总结

算子类型	数据复用度	计算密度	缓存敏感度	融合收益
Conv	高（权重复用）	中高	高	高（减少中间结果写回）
GEMM	中（分块后高）	高	极高	中（主要靠分块优化）
Element-wise	低（无复用）	极低	低	极高（消除冗余访存）

从这张表能看出来，逐元素算子的融合收益最高，因为它本身不做什么计算，纯粹在搬运数据。卷积和GEMM的融合收益更多体现在减少中间结果的写回上。

知识体系结构图

下面这张图，我把三种算子的数据流和融合关系画出来了，方便你理解：

从这张图能看出来，三种算子的数据访问模式差异很大。卷积靠权重复用，GEMM靠分块复用，逐元素算子几乎不复用。融合时，我们要把「高复用」的算子放在前面，把「低复用」的算子跟在后面，让数据在缓存里多待一会儿。

嗯，这就是我这些年总结的经验。下次你看到网络结构里有一串 Conv + ReLU + BatchNorm，别犹豫，直接考虑融合。收益往往超出你的预期。

想要更多资料，联系可可