深入浅出 tvm – (11) Relay Pass 之算子融合

算子融合,FuseOps,是编译领域最常见的一个优化技术,在这里也算属于 relay 里面最复杂的一类优化了,整个优化的核心逻辑 1k+ 行
代码:src/relay/transforms/fuse_ops.cc
算子融合的目的最终是要解决 AI 处理器的内存墙、并行墙的问题,提升 Tensor 数据的访存局部性。目前算子融合的技术路线有比较多,这里不涉及,我们只需要知道 tvm 是基于支配树来实现算子融合的就行了

1. 基本概念

1.1. 算子融合

算子融合,即将多个算子组合在一起放到同一个核中,通过算子融合的方式,不需要将中间结果保存到全局内存,进而减少执行所需要的时间。
tvm 中将算子分为7种类型:
  1. kElemWise:2个 tensor 之间按元素逐个操作的算子,实际上所有四则运算都是这种类型,https://deeplizard.com/learn/video/QscEWm0QTRY
  2. kBroadcast:见上述链接,到操作两个不同形状的 tensor 时
  3. kInjective:一对一映射函数,比如 add / sqrt / exp 等操作算子(operator)
  4. kCommReduce:多到少的映射,输入到输出具有降维性质,如:sum / max / min等操作操作算子(operator)
  5. kOutEWiseFusable:这是计算比较复杂的,如:conv2d / bn /  relu等操作算子(operator)
  6. kTuple:xx
  7. kOpaque:无法被融合的算符,比如 sort
根据以上对算符的不同类型,TVM提供了三种融合规则(我看论文是这么写的,但是现在不止3种了?):
0
从融合算子的内部视角看,这种融合实际上是数据计算pipeline化,即两次计算中间数据不再经历store-load过程,而是直接给到下一个计算单元完成计算。