TVM图级优化了解

发布时间 2023-03-27 00:21:18作者: 牛犁heart

TVM图级优化按照优化范围,可分为局部优化全局优化
局部优化是TVM图级优化的重点,其中算子融合是AI编译器必不可少的优化方法。
算子融合核心思想就是将多个算子合并成一个内核,因而无需将中间结果写回全局内存,减少了中间变量的分配,也减少了片上缓存和片外存储之间的数据传输。

那算子融合也是有条件的,不能各算子谁想和谁融合就能融合的

不同算子能融合在一起的前提条件就是:某个算子访问的数据是其他算子在其共享存储资源内生成的数据,即算子间应满足访存的局部性

好处有哪些呢?

通过算子融合(如合并循环嵌套)可以为进一步优化创造条件,还可以减少启动和同步开销,更好的共享计算资源。
虽然就单个算子而言,融合后的算子可能会增加存储资源的开销,但考虑到整体性能的提升,这中开销是值得的。

不好的地方呢?

融合后的算子可能破坏后续算子之间的并行执行,即在融合前,某个算子的后续算子是并行执行的,但在将该算子和其他算子融合后有些后续算子只能在融合算子执行结束后开始执行,这显然会影响性能。

局部优化还有什么?

局部优化中还包括经典的代数简化。代数简化考虑的是一系列节点,利用不同类型节点的交换律、结合律和分配律来简化计算。如强度消减、常量折叠等。

  • 强度消减:可以用开销小的算子代替开销大的算子。
  • 常数折叠:通过常数折叠,可以预先计算出可静态确定的图,用常量值替换常量表达式,从而节省执行开销。

代数简化也可以应用于AI相关算子(如reshape、转置和池化等)

全局优化又是什么呢?

计算图提供了全局视图,TVM可以在整个计算图中搜索特定特征,并针对这些特征执行特定全局优化操作。
例如:计算图中的死代码通常是由其他图优化造成的。因此,在其他图优化之后还应执行死代码清除(Deal Code Elimination,DCE)公共子表达式消除(Common Subexpression Elimination,CSE)操作