算子融合,FuseOps,是编译领域最常见的一个优化技术,在这里也算属于 relay 里面最复杂的一类优化了,整个优化的核心逻辑 1k+ 行
代码:src/relay/transforms/fuse_ops.cc
算子融合的目的最终是要解决 AI 处理器的内存墙、并行墙的问题,提升 Tensor 数据的访存局部性。目前算子融合的技术路线有比较多,这里不涉及,我们只需要知道 tvm 是基于支配树来实现算子融合的就行了
1. 基本概念
1.1. 算子融合
算子融合,即将多个算子组合在一起放到同一个核中,通过算子融合的方式,不需要将中间结果保存到全局内存,进而减少执行所需要的时间。
tvm 中将算子分为7种类型:
- kElemWise:2个 tensor 之间按元素逐个操作的算子,实际上所有四则运算都是这种类型,https://deeplizard.com/learn/video/QscEWm0QTRY
- kBroadcast:见上述链接,到操作两个不同形状的 tensor 时
- kInjective:一对一映射函数,比如 add / sqrt / exp 等操作算子(operator)
- kCommReduce:多到少的映射,输入到输出具有降维性质,如:sum / max / min等操作操作算子(operator)
- kOutEWiseFusable:这是计算比较复杂的,如:conv2d / bn / relu等操作算子(operator)
- kTuple:xx
- kOpaque:无法被融合的算符,比如 sort
根据以上对算符的不同类型,TVM提供了三种融合规则(我看论文是这么写的,但是现在不止3种了?):
从融合算子的内部视角看,这种融合实际上是数据计算pipeline化,即两次计算中间数据不再经历store-load过程,而是直接给到下一个计算单元完成计算。