图优化技术

常见的图优化技术包括常量折叠、公共子表达式消除、代数化简、算子融合等技术，接下来将分别简单介绍其原理。

1.常量折叠 (Const Folding)

常量折叠是传统编译器中的一种优化手段, 它的原理是如果一个计算所依赖的所有输入都是常量, 则在编译期间就可以得到计算结果. 比如对于下面的代码:

int main() {
    const int a = 100;
    int b = a + 100;
    printf(“b is %d\n”, b);
    return 0;
}

b 的值只依赖于 a 的值, 由于 a 为常量, 其取值在编译阶段就确定了, 因此也可以在编译阶段计算得到 b 的值, 所以上面的代码经过常量折叠优化后,等价于下面的代码:

int main() {
    printf(“b is %d\n”, 200);
    return 0;
}

深度学习编译器中的常量折叠和传统编译器是类似的, 只需将输入变为 Tensor 即可. 比如对于下面的一个网络:

A = const([3, 5], 1.0, fp32)
B = const([3, 5], 0.5, fp32)
C = var([5, 4], fp32)
D = dot(A + B, C)

通过常量折叠后, 就可以变为:

TMP = const([3, 5], 1.5, fp32)
C = var([5, 4], fp32)
D = dot(TMP, C)

2. 公共子表达式消除 (Common Sub-Expression Elimination)

在一个程序中, 如果几个表达式的类型、参数和输入均相同, 则将他们称做公共子表达式。 对于公共子表达式, 只需要计算其中一个表达式的值, 其他表达式的值可以通过赋值得到。这个过程就称作公共子表达式消除, 它是一种传统编译器中常用的优化手段, 经过迁移也可以应用到深度学习编译器中。

基本思路是通过一个 MAP 表, 记录截止当前, 已处理过的同一种类型的 OP。 对于当前正在处理的 OP, 先查找该 MAP 表, 如果能找到其他和正在处理的 OP 类型相同的 OP, 则对他们进行遍历, 如果其中某个 OP 的输入和参数与当前正在处理的 OP 相同, 则它们为公共子表达式, 结果可以互相替代；如果所有 OP 都不能与当前正在处理的 OP 匹配, 则将当前 OP 复制一份返回。

3. 代数化简

代数化简的目的是利用交换率、结合律等规律调整图中算子的执行顺序，或者删除不必要的算子，以提高图整体的计算效率。

比如，在下面的例子中，sub_module_1 在输出前将结果 Reshape 为三维 Tensor；而 sub_module_2 要求输入是二维 Tensor，因此一开始就会对输入 Tensor 进行一次 Reshape 操作。将sub_module_1 和 sub_module_2 级联起来后网络中会有两次连续的 Reshape 操作，但很明显这里只需要做一次 Reshape 即可。

通常有两种场景会导致出现例子所示情况。第一种就是例子中所示的封装，随着网络变的越来越复杂，不可避免的要调用其他已经封装好的子网络，为了适配子网络的接口，就会加入 Reshape、Transpose 等操作。第二种是编译器在图优化过程中插入了新的算子，而插入的新的算子和已有的算则之间存在冗余。

def sub_module_1():
    ...
    ret = ...
    return tf.reshape(ret, [-1, n, k])

def sub_module_2(inputs):
    i0 = tf.reshape(inputs[0], [-1, k])
    xxxx

def main():
    s1 = sub_module_1()
    s2 = sub_module_2([s1])
    xxx

下表列举了一些 Tensorflow 中支持的代数化简操作：

4. 算子融合

算子融合的目的是将几个小 OP 融合为一个大 OP，达到减少从内存/显存中搬移数据的目的。

举例来说，假设要计算 Relu(X + Y），X 和 Y 的长度均为 L。在 Tensorflow 中会对应到 2 个 OP：Add 和 Relu。做 Add 计算时，先要将 X 和 Y 从内存/显存中读出，然后再将计算结果写入到内存/显存，因此 Add 计算需要从内存中读写的数据量为 3L。做 Relu 计算时，同样需要将输入数据从内存/显存中读出，然后再将结果写回到内存/显存，因此 Relu 计算需要从内存中读写的数据量为 2L。Add 和 Relu 加起来读写数据的总量为 5L 。如果我们将其融合为一个 OP，将 X 和 Y 从内存/显存中读出后，先做加法，然后做 Relu 计算，再将最后的结果保存到内存/显存，这样读写数据的总量仅为 3L，相比融合前减少 40%。

下表列举了几个 Tensorflow 中支持的算子融合：

Conv2D + BiasAdd + <Activation>
Conv2D + FusedBatchNorm + <Activation>
Conv2d + Squeeze + BiasAdd
Matmul + BiasAdd + <Activation>

注：<Activation> 表示该 OP 是可选项。

从上面的例子可以看出，算子融合有效果需要满足几个前提条件。

首先，中间结果只能写回到最后一级存储。如果 Add 的计算结果可以保存到高速缓存中，读取速度非常快，那融合前读写内存/显存的数据量也可以近似认为是 3L，而非 5L。

其次，最终结果只依赖于输入数据的部分数据，而非所有数据，否则没有办法在减少访存的情况下实现融合后的计算逻辑。

最后，从内存/显存中读写数据的耗时占整个 OP 计算过程中的占比较高，只有满足这个条件，减少数据搬运量才能起到加快图运行的效果。

和代数化简等其他图优化手段相比，算子融合有一个很大的不同：引入了新的融合后的 OP。因此算子融合的难点在与如何为融合后的 OP 实现计算逻辑？

目前有两种模式：

一种是将融合后的算子做为普通算子, 为他们实现定制 Kernel; PaddlePaddle, Tensorflow 等框架本身采用的是这种方式。 但这种模式开发成本较高，相对来说能够支持的融合模式有限。

还有一种模式是和自动代码生成结合, 直接生成可以运行的代码。 TVM 和 XLA 采用的是这种方式。