图层优化

将一种计算图结构，在不改变算数结果的情况下，基于设定好的规则，对计算图进行相应的图替换操作．

  • 读写冗余：一些计算场景中存在重复读写内存、或者内存访问不连续，降低 cache hit rate，导致多余的内存传输
  • 结构冗余：模型存在无效的计算节点、重复的计算子图、相同的结构与模块

图层优化包括：

  1. 常量折叠
  2. 冗余节点消除
  3. 算子融合
  4. 数据布局转换

常量折叠

下面计算图（左图）里，A,B,CA,B,C 这个子图中，由于 A,BA,B 都是常量，其计算结果也是常量，那么我们可以预先计算好 A,B,CA,B,C 这个子图的结果，然后直接“折叠”为 Const 3\texttt{Const 3}，结果如右图所示：

graph TB;subgraph foldable [foldable];A["Const 1"] --> C["Op 1"];B["Const 2"] --> C;end;C --> E["Op 2"];D["Input 1"] --> E;
graph TB;A(["Const 3"]) --> C["Op 2"];B("Input 1") --> C;

除此之外，有一些计算节点需要的是张量的 shape 数据。当张量的形状可以在编译器确定（即便张量具体的值是未知的），AI 编译器也可以对获取 shape 这一节点进行优化。

冗余节点消除

AI 编译器中，这一环节的优化类似于传统编译器里的死代码消除 (Dead Code Elimination)

节点本身无意义

graph TB;
A("Input A") --> B[Concat]
B --> D[Subtract]
C("Input B") --> D

这个 concat 算子只有一个参数，相当于没有用，因此属于“本身无意义”的节点，因此可以直接消除：

graph TB;
A("Input A") --> C[Subtract]
B("Input B") --> C

公共表达式消除

出现公共计算子图时，可以将这个公共的子图进行复用。

算子融合

将多个细粒度的计算操作合并为单个符合算子，减少内核启动次数和中间结果的频繁 IO，有效降低计算开销和内存带宽压力．

常见的优化思路包括：

  • 将多个算子融合为一个已知的算子
  • 融合成自定义的大算子
  • 融合到算子的某一个属性

卷积 conv 与 batchnorm 进行融合

卷积的数学操作是

yi=Wconv∗xi+Bconvy_i=W_{conv}\ast x_i + B_{conv}

而 BatchNorm 的操作是

yi=γxi−μσ2+ϵ+β=γxiσ2+ϵ+(β−γμσ2+ϵ)\begin{aligned} y_i&=\gamma \frac{x_i-\mu}{\sqrt{\sigma^2+\epsilon}}+\beta \\ &=\gamma \frac{x_i}{\sqrt{\sigma^2+\epsilon}}+(\beta-\frac{\gamma\mu}{\sqrt{\sigma^2+\epsilon}}) \end{aligned}

所以我们可以结合 Conv 和 BatchNorm

yi=γσ2+ϵ(Wconv∗xi+Bconv)+(β−γμσ2+ϵ)\begin{aligned} y_i=\frac{\gamma}{\sqrt{\sigma^2+\epsilon}}(W_{conv}\ast x_i + B_{conv})+(\beta-\frac{\gamma\mu}{\sqrt{\sigma^2+\epsilon}}) \end{aligned}

数据布局转换

例如，对于图像模型来说，通常有 N,C,H,W 和 N,H,W,C 两种数据布局．数据布局主要影响的是算子计算中的 cache hit rate

一般来说，我们可以在算子的前后插入 Transpose 算子，从而适配底层硬件平台．