3. 算子优化基础：常见算子的计算特性与优化思路

各位同学，今天我们来聊聊算子优化。这是部署调优里最基础、也最见功力的一环。

我刚开始做部署时，总觉得算子就是调个库，有什么好优化的？直到有一次，一个简单的 Conv 层在手机端跑了 200ms，怎么调都下不来。后来才发现，是内存访问模式出了问题。嗯，从那以后，我再也不敢小看算子优化了。

3.1 算子的两大分类：访存密集型 vs 计算密集型

要优化算子，首先得搞清楚它「缺什么」。说白了，就是看瓶颈在计算还是在访存。

计算密集型算子：算力是瓶颈。GPU 的 ALU 在疯狂干活，但显存带宽还有富余。典型代表：大尺寸的 Conv、MatMul。

访存密集型算子：带宽是瓶颈。ALU 大部分时间在等数据，计算单元吃不饱。典型代表：Activation（ReLU、Sigmoid）、Element-wise 操作、BatchNorm。

怎么区分？我有个简单的经验：看计算密度（Arithmetic Intensity）。计算量（FLOPs）除以访存量（Bytes），比值越大，越偏向计算密集型。

核心判断标准：

  • 计算量 / 访存量 > 100：计算密集型
  • 计算量 / 访存量 < 10：访存密集型
  • 中间地带：两者都要考虑

你想想看，一个 3x3 的 Conv，如果输入通道很少，那计算密度其实不高，反而可能是访存瓶颈。我在项目中就踩过这个坑，以为 Conv 都是计算密集，结果优化方向全错了。

3.2 Conv 算子：计算特性与优化思路

Conv 是深度学习里最核心的算子，没有之一。它的计算特性很明确：高计算量、高数据复用。

一个典型的 Conv 层，计算量是：

FLOPs = 2 × H_out × W_out × C_out × C_in × K_h × K_w

访存量呢？输入特征图、权重、输出特征图，加起来也不小。但关键是，权重可以被多个输出位置复用，输入特征图也可以被多个卷积核复用。

优化思路，我总结了几条：

  1. Im2Col + GEMM：把 Conv 转成矩阵乘法。这是最经典的做法，利用高度优化的 BLAS 库。缺点是需要额外的内存来存储展开后的矩阵。
  2. Winograd 算法：对小卷积核（3x3）特别有效。我记得在某个项目中，用 Winograd 把 3x3 Conv 加速了 1.5 倍。但要注意，它对精度有一定影响。
  3. 内存布局优化：NCHW 还是 NHWC？我个人习惯用 NHWC，因为它在很多硬件上访存更连续。尤其是 ARM 平台，NHWC 配合内联汇编，效果拔群。
  4. 算子融合：把 Conv + BN + ReLU 融合成一个算子。减少中间结果的读写，访存开销能降 30% 以上。

我的经验：对于 1x1 Conv，它本质上就是矩阵乘法。直接用 MatMul 的优化策略，比用 Conv 的优化策略更高效。别问我怎么知道的，调了三天才发现。

3.3 MatMul 算子：计算特性与优化思路

MatMul（矩阵乘法）是另一个计算密集型算子。Transformer 火了之后，它的地位直线上升。

计算特性：O(n³) 的计算量，O(n²) 的访存量。n 越大，计算密度越高，越偏向计算密集型。

优化思路：

  • 分块（Tiling）：把大矩阵切成小块，让数据尽量留在 L1/L2 cache 里。我一般把块大小设为 64x64 或 128x128，具体要看硬件 cache 大小。
  • 内存对齐：确保矩阵的行首地址是 16 字节或 32 字节对齐。不对齐的话，SIMD 指令效率会暴跌。我曾经因为一个结构体没加 alignas，性能直接腰斩。
  • 使用 FMA 指令：乘加指令，一条指令完成乘法和加法。现代 CPU 和 GPU 都支持，能有效提升计算密度。
  • 稀疏化：对于大模型，可以用稀疏矩阵乘法。但要注意，稀疏度不够高时，反而更慢。我个人建议稀疏度低于 70% 就别用了。
// 一个简单的分块矩阵乘法示例（伪代码）
for (int i = 0; i < M; i += BLOCK_SIZE) {
  for (int j = 0; j < N; j += BLOCK_SIZE) {
    for (int k = 0; k < K; k += BLOCK_SIZE) {
      // 加载 A[i:i+BLOCK][k:k+BLOCK] 到 L1 cache
      // 加载 B[k:k+BLOCK][j:j+BLOCK] 到 L1 cache
      // 计算 C[i:i+BLOCK][j:j+BLOCK] += A_block * B_block
    }
  }
}

3.4 Activation 算子：计算特性与优化思路

Activation 是典型的访存密集型算子。ReLU、Sigmoid、Tanh，计算量很小，但需要读写整个特征图。

计算特性：计算量极低，访存量巨大。一个 ReLU 就是一个比较指令，但数据要从显存搬到寄存器，算完再写回去。这个搬运过程才是瓶颈。

优化思路：

  • 算子融合：把 Activation 融合到前一个算子（Conv、MatMul）里。这是最有效的优化，没有之一。融合后，数据不用写回显存，直接在寄存器里完成。
  • 使用查表法：对于 Sigmoid、Tanh 这类复杂函数，可以用查表代替计算。精度要求不高时，效果很好。我试过用 256 个元素的查找表，精度损失不到 0.1%，速度提升了 3 倍。
  • 向量化：用 SIMD 指令一次处理多个元素。ReLU 可以用一条指令完成 8 个 float 的比较。
  • 减少数据搬运：如果实在不能融合，尽量让 Activation 在数据还在 cache 里时执行。比如，在 Conv 的每个输出通道计算完后，立即执行 ReLU。

避坑指南：我曾经把 ReLU 单独写成一个 kernel，结果发现整个模型 40% 的时间花在 ReLU 上。后来融合到 Conv 里，时间直接降到 5%。记住：访存密集型算子的优化核心是「少搬数据」。

3.5 如何判断一个算子的类型？

我一般用 Roofline Model 来分析。简单说，就是画一条天花板曲线：

  • 横轴：计算密度（FLOPs/Byte）
  • 纵轴：性能（FLOPs/s）
  • 斜线部分：访存瓶颈
  • 水平部分：计算瓶颈

把算子的计算密度画上去，落在斜线区域就是访存密集型，落在水平区域就是计算密集型。

举个例子：

算子 计算量 (FLOPs) 访存量 (Bytes) 计算密度 类型
Conv 3x3 (256x256x64) 1.2e9 8.4e6 143 计算密集型
ReLU (1e6 元素) 1e6 4e6 0.25 访存密集型
MatMul (1024x1024) 2.1e9 8.4e6 250 计算密集型

你看，ReLU 的计算密度只有 0.25，明显是访存瓶颈。而 Conv 和 MatMul 的计算密度都在 100 以上，是计算瓶颈。

我的习惯：拿到一个新模型，先跑一遍 profiling，看看每个算子的耗时占比。然后按计算密度排序，优先优化那些「耗时高且计算密度低」的算子。这些往往是优化收益最大的地方。

3.6 总结

算子优化，说白了就是搞清楚「瓶颈在哪」。计算密集型就优化计算（分块、SIMD、FMA），访存密集型就优化访存（融合、减少搬运、查表）。

我见过太多人一上来就调 kernel，结果方向错了，白费功夫。记住：先分析，再优化。用 Roofline Model 画个图，一目了然。

下一章，我们会深入讲 Conv 的 Winograd 算法实现。到时候我会分享一个实际项目中的踩坑经历，保证让你少走弯路。

← 上一章 📖 返回目录 下一章 →