图替换 | Don't Respond

Contents

1. 背景
2. 经典图替换
3. 基于超优化的图替换
4. 低精度和layout优化
5. 相关资料

背景

图替换（或者叫图改写）是一种重要的图优化技术，几乎在所有的开源框架（尤其是移动端框架）中都有应用。比如tensorflow r1.14版本中就包含了155个替换子，而且实现这些替换子的总代码量接近53k行。

一些常见的图优化技术：

DCE

CSE（公共子表达式消除）

常量折叠

数学公式简化

Op融合

Layout变换

内存优化（swap-in/swap-out、重计算）

由于目前的编译器技术通常基于low-level的中间表达，注重对局部计算的优化，对于跨多个粗粒度op的优化要不无能为力，要不就得增加编译器的分析难度并导致代码膨胀。一般来说AI框架支持的粗粒度op非常有限，而且这些op的组合常常也比较固定，比如convolution通常和bias_add、relu组合使用，因此基于高层中间表达的图替换成为一种比较可行的优化方案。经过图替换优化后的计算图再经过编译器的优化后，生成最终的硬件代码。

目前主流开源框架的图替换都是基于经验和手工设置的替换子来实现的，在这里统称为经典图替换技术。

经典图替换

图替换是将原始计算图替换成另一个优化后的等价计算图，替换后的计算图通常是硬件友好的，比如可以消除中间结果，降低内存占用，减少访存和计算量，并且不影响最终的计算结果。

在进行图替换之前，首先需要定义出源计算图到目标计算图的替换规则（替换子），由于这些替换规则往往需要依靠人的经验一条条手工去定义，因此称之为经典图替换。给出一条替换子，我们需要在原始计算图中不断地去匹配替换子的源计算子图，一旦匹配到满足要求的子图后，就将源计算子图重新映射为替换子中的目标计算图。

在一些开源框架中，替换子的定义形式不尽相同。在TensorFlow中源图匹配和替换的定义是非常松散的，它甚至没有直接定义出替换子的源图，而是定义一系列约束来判断是否匹配。PaddlePaddle中则是将一个替换过程定义为一个pass，pass执行时动态构建相应的替换子源图，执行匹配算法并回调源图到目标图的替换函数。比如下面是TensorFlow中将Conv+BiasAdd替换成FusedConv的过程。

定义匹配约束

struct ContractionWithBiasAdd {
  int constraction;
  int bias_add;
}
// node为输入的grapper node， pattern为输出的ContractionWithBiasAdd.
bool FindContractionWithBias(node，*pattern) {
  // 开始列举匹配的constractions.
  1、如果node存在控制边，返回false
  2、如果node不是BiasAdd，返回false
  3、如果node的父节点不是Conv或MatMul，返回false
  4、...
  // 如果以上所有constructions都满足，则将需要替换的node id写到特定的pattern中。
  pattern->constraction = node的父节点;
  pattern->bias_add = node；
  return true;
}

定义替换过程

// pattern为输入的ContractionWithBiasAdd，
void AddFusedContractionNode(pattern, *invalidated_nodes) {
  1、创建一个新的node：fused_op
  2、将Conv或MatMul的input和filter添加到fused_op的输入中，并将BiasAdd的bias加到fused_op的输入
  3、根据Conv或MatMul的一些参数设置fused_op的参数，比如conv的kernel、channel、padding等等，以及matmul的transpose等
  4、将fused_op加入到graph，同时将Conv或MatMul和BiasAdd加入到invalidated_nodes
}

TensorFlow采用的定义匹配约束的方式与直接定义出子图的方式本质上是等价的，但相比后者可读性较差，而优点就是代码可复用性高，比如上面的FindContractionWithBias可以同时匹配Conv+BiasAdd和MatMul+BiasAdd两种子图，并且这些约束便于嵌套使用。

无论是TensorFlow还是PaddlePaddle，图替换都是不完全的。比如说对于Conv+BiasAdd+BiasAdd这种计算图，第一次只能匹配到Conv+BiasAdd，替换后又变成了一个Conv+BiasAdd的计算图，因此TensorFlow中默认采用了两遍优化。根据TensorFlow公开的一些数据，基本上第二次优化的机会已经非常少了。

InceptionV3
Seq2Seq

基于超优化的图替换

超优化（Superoptimization）是现代编译器中的一种指令优化技术，其主要工作原理是通过随机生成指令序列以及暴力搜索的方式自动找到一组优化的指令序列，并等价替换原有的指令序列。1992年第一个Superoptimizer被集成到了GCC编译器，之后Google也为LLVM开发了一个Superoptimizer，取名为Souper。

依靠人工设定的编译器往往对代码的优化不够彻底，给生成的code留下了大量的优化空隙，而且人工设定的优化规则往往没有经过充分验证，经常导致各种极端条件下的代码bug。Superoptimization将指令序列优化问题转换为自动搜索问题，并加入了自动化验证和一阶逻辑验证，在发现代码优化空隙的同时优化结果也更加可靠。

TASO（Tensor Algebra SuperOptimizer）将Superoptimization用于DNN高层中间表达的图优化，在大多数模型上取得了比XLA和TensorRT更优的效果。TASO的工作是MetaFlow（作者另一个基于人工规则的图替换框架）的延续，因此也采用了与MetaFlow一致的替换子定义。MetaFlow替换子的定义包括：源图、目标图、输入和输出的映射关系。

TASO相比其他开源框架最大的区别就是不需要手工去设定各种各样的替换子，只需要像设计硬件指令一样设计出基本的算子定义（或者计算逻辑），之后系统会根据指定的算子集自动生成满足条件的替换子，经过验证的替换子最终作用于图替换过程。基于高度抽象的替换子定义，TASO可以独立于具体的训练或预测框架，离线完成替换子的生成和验证，并在图优化阶段加载到程序中进行图替换。尽管手工设计有很多弊端，但TASO在代码实现过程中并没有完全抛弃手工设计的方式，而是采用了手工设计和替换子自动生成相结合的方式。

替换子定义

替换子包含三个部分，源图、目标图、输入和输出tensor的映射关系。并且替换子通常是与shape无关的，源图和目标图都是由算子构成的，每个算子都可以指定一些配置，比如kernel指定卷积核的大小、axis指定reduce的维度等等。

但需要注意的是concat和split两个算子，在图替换中这两个算子通常用于算子融合，比如下图对两个不同的输入B和C进行相同的MatMul操作，就可以替换为先将输入B和C进行一次合并，然后调用一次MatMul后，对结果进行切分得到两个输出X和Y。

为了能正确切分出X和Y，在Concat时我们需要给每个维度维护一个分割树（split tree）。一个行分割的例子如下，图中需要将A和B按照第0维进行concat，因此输入A在第0维有一个原始的分割树[0, \(S_{A}\)]，表示对于tensor A，第0维从0到\(S_{A}\)行都是A的数据区域。A和B concat后tensor的row变成了\(S_{A}+S_{B}\)，并且通过分割树可以知道第0到\(S_{A}\)行是A的数据，从\(S_{A}\)到\(S_{A}+S_{B}\)行是B的数据。根据分割树，Split非常容易地就可以将数据进行切分。TASO的分割树支持任意维度的切分和递归切分。

替换子生成

替换子生成包含两个阶段：构建搜索空间，以及对潜在的替换子进行测试。

构建搜索空间

搜索空间由任意合法的计算图构成，而计算图由给定的算子集中的算子组成。TASO向我们表明了一种暴力枚举、深度优先递归构建的方法。

给定算子集和初始化的input tensor集合，对于每一个输入tensor，每次从算子集中选择一个合法的算子构建graph，并计算当前graph的输出tensor，将输出tensor加入到input tensor集合，保存graph以及graph的fingerprint（对输出tensor计算hash值），接着重复上面的过程继续加入算子，直到递归的深度达到设定的上限。

对于同样的输入tensor，如果构建的两个计算图的输出tensor相同，则这两个计算图构成了一个潜在的替换子。为了避免出现浮点计算异常的情况，构建计算图时所有的tensor都是int类型。
测试潜在替换子

为了进一步验证潜在替换子的合法性，TASO设计了一系列cases来测试潜在替换子。每个测试case都使用随机初始化的输入tensor，当两个计算图结果一致时才认为测试通过，只有所有测试cases都通过的潜在替换子才是合法的替换子。

与构建计算图时使用int类型的tensor不一样，所有测试case的输入tensor都是-1到1之间的浮点数。由于relu对于所有小于0的值都返回0，因此可能导致非法的替换子通过测试cases，作者认为可以使用任意一个非线性函数来代替relu，TASO中使用\(x(x+1)+1\)。

替换子验证

TASO同时使用一阶逻辑表达的算子属性对替换子进行进一步验证，这些属性通常是由人工定义，并且经过充分review和大量测试验证过的。

在定义算子属性之前，首先需要对算子进行符号建模，算子模型通常包含参数和输入tensors。比如\(conv(s, p, c, x, y)\)表示conv算子的符号模型，\(s\)，\(p\)，\(c\)是conv的参数，分别表示stride、padding和activation，\(x\)和\(y\)是卷积操作的两个输入。如果activation是none，很显然conv就是一个线性操作，因此满足以下属性： \[ \begin{aligned} ∀s,p,x,y,z. conv(s,p,Anone,ewadd(x,y),z) = \\ ewadd(conv(s,p,Anone,x,z),conv(s,p,Anone,y,z)) \end{aligned} \] TASO定义了大量的算子属性，并且使用z3（一阶逻辑验证器）对所有合法的替换子进行验证。如果有合法的替换子无法被一阶逻辑验证，则需要根据替换子手动添加一条算子属性，以确保所有合法的替换子都能验证通过。

冗余替换子裁剪

自动生成的替换子往往存在大量的冗余，TASO使用了两种策略消除冗余。

Input tensor renaming

对输入进行重命名的方式消除不同替换子之间的冗余。比如下面两个替换子，

替换子a:

替换子b：

将替换子a的一个输入tensor A改名为C，就得到了替换子b，说明这两个替换子存在冗余，因此最终只会保留更加通用的替换子b。
Common subgraph

如果替换子的源图和目标图包含同样的子图，则可以用一个相同的tensor替换掉公共子图。比如下面的一个替换子，

source graph和target graph包含同一个子图（B x C），将source graph替换成target graph时，公共子图没有任何变化，因此可以将子图消除。