在专栏CUDNN中，可以看到CUDNN库提供了许多算子，并且有着大量的优化，这里面主要的方法就是算子融合，cuda编写了许多高性能的算子/融合算子，以及许多Engines，包括预编译/运行时。本文讨论一下算子融合的原理与一些开源实现方式。
算子融合介绍
算子融合（Operator Fusion）是一种优化技术，用于将多个连续的独立算子（operations）合并成更大的算子，以减少计算和内存访问的开销，提高计算效率和性能。算子融合通常在深度学习框架中应用，有助于优化计算图并提高模型执行的效率。
对于上图所示的一个简单的网络结构，写成公式可以表示为： <math xmlns="http://www.w3.org/1998/Math/MathML"> y = a ( x + b ) + x y=a(x+b)+x </math>y=a(x+b)+x
看上去似乎很简单，但是对于底层硬件来说，由于是流式处理数据，步骤就可以解析成：
<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"> y = x + b y = a ∗ y y = y + x y=x+b\\ y=a*y \\ y=y+x </math>y=x+by=a∗yy=y+x
如果不使用算子融合，每一个步骤都是把数据从DDR搬运到专用的处理器(CPU的话就是cache，寄存器）上，进行运算后，再将结果返回到DDR上的。所以，如果要一次性将这三个步骤在一个cycle中全部算完再吐出的话，首先要存张量x， 要存张量a， b，要存中间量或者说最终输出量y，而这个y你还不能复用x的空间，因为等下还要加x，这对专用硬件的要求是挺高的，同时对于框架也要求有所提升。
而通过算子融合的方式，可以减少中间结果的存储和传输，有效的减少访存的开销；另外，合并多个算子可以减少计算的次数，在GPU等并行计算设备上，可以有效提高计算效率。
算子融合主要为了解决内存墙和并行墙两个问题。
内存墙
算子融合技术解决内存墙问题，在不同框架上呈现出两种典型的发展路线：
- 遍历路线：针对特定设备，枚举实现典型的CB（计算密集型，Compute-bound）+ MB（访存密集型，Memory-bound）形式的融合算子，如Conv+ReLU/Conv+BN+ReLU/Dense+ReLU/Conv+Sum等等，Intel的oneDNN以及国内很多大厂的推理框架走的都是这个路线。
- 规则路线：基于规则实现算子融合，例如XLA、TVM、MLIR。以TVM为例，其将所有算子分为Opaque/Injective/Reduction/Complex-out四类，并设定融合规则形如Complex-out+Injective/Injective+Injective/Injective+Reduction等，以此实现所有符合规则的子图融合。
并行墙
解决并行墙问题，即如何保证计算图执行非CB部分的子图时也能充分利用设备的计算资源。这方面典型的工作是TASO和RAMMER，通过将计算过程中互不依赖的低计算量子图打包为一个kernel并行计算以提升资源利用率。RAMMER提出一种将计算图中的算子节点进行并行编排，从而提升整体计算并行度。特别是对于网络中存在可并行的分支节点，这种方式可以获得较好的并行加速效果。
XLA
XLA（Accelerated Linear Algebra），是一个深度学习编译器。
在此介绍一下Tensorflow的基本运行模式。对于Tensorflow来说，其有内置的算子库，搭建一个模型，最终成果就是构建了一个由算子组成的图。
Tensorflow的运行时（Executor）加载所有的图依赖的算子（Kernels），从计算图的根节点开始跑，遇到算子就把kernels里对应的实现调出来执行，然后根据节点间的依赖关系，判断哪些节点现在依赖满足可以调度，哪些还无法调度，来调度相应的节点。
所以，如果把图描述看成是高级语言，上图有点解释执行的意思。那么，将其改为编译执行就是提高性能的方法，而分层的方法在编译领域很有用。
具体而言，得先给Tensorflow的图描述转成一种XLA的中间表达形式，然后在往后进行编译优化。 先转成XLA的图表达之后，还要再往下层转，整个路径是Tensorflow->XLA->LLVM（开源编译框架）->各种硬件系统的汇编->机器码，如下图所示。
在XLA层可以做自动的算子融合，这部分通常被称作HLO（high-level optimizer），如下图所示。主要就是通过算子融合省通信和访存开销，复杂场景下，通信和访存永远是大头。
运行模式
即时编译（JIT）
这是一种在对用户生产流程侵入较小的情况下使得编译生效的方式，对训练和推理过程都有效。
不是所有的TF算子都能够被翻译成XLA的中间表达。那么在图被调度运行时的时候，需要先识别出来在TF计算图中，哪部分可以被转化，哪些不行。能够成功转化的子图被圈选出来，这个过程称为聚类（Clustering）。后续TF运行时在进入子图之前，就会触发JIT编译过程，直接调用编译后的子图代码并缓存起来下次继续用。（计算图中节点大部分都会被多次运行，重复调用、迭代式的计算训练等等）
对于一个复杂的XLA Graph，需要在有限的运算开销下避免死锁等情况发生，需要采用一些贪心、启发式的算法，实际挺复杂的，这部分需要要看相关理论和源码。
提前编译（AOP）
更多应用场景在推理优化阶段，针对整张图提前进行编译。（当然这要求整个图中的算子都是可以被转化为XLA的中间表达的）。
XLA工作示例
我们先来看XLA如何作用于计算图，下面是一张简单的计算图
这里假设XLA仅支持matmul和add。XLA通过图优化方法，在计算图中找到适合被JIT编译的区域
XLA把这个区域定义为一个Cluster，作为一个独立的JIT编译单元，计算图中通过Node Attribute标示
然后另一个的图优化方法，把cluster转化成TensorFlow的一个Function子图。在原图上用一个Caller节点表示这个Function在原图的位置
最后调用TensorFlow的图优化方法(BuildXlaOps)，把Function节点转化成特殊的Xla节点。
在TensorFlow运行时，运行到XlaCompile时，编译Xla cluster子图，然后把编译完的Executable可执行文件通过XlaExecutableClosure传给XlaRun运行。
接着根据虚拟指令分配GPU Stream和显存，然后IrEmitter把HLO Graph转化成由编译器的中间表达LLVM IR表示的GPU Kernel。最后由LLVM生成nvPTX（Nvidia定义的虚拟底层指令表达形式）表达，进而由NVCC生成CuBin可执行代码。
TVM
TVM是一个基于编译优化的深度学习框架，TVM的架构图如下图所示：
从这个图中我们可以看到，TVM架构的核心部分就是NNVM编译器（Relay也可以看作是NNVM）。NNVM编译器支持直接接收深度学习框架的模型，如TensorFlow/Pytorch/Caffe/MxNet等，同时也支持一些模型的中间格式如ONNX、CoreML。这些模型被NNVM直接编译成Graph IR，然后这些Graph IR被再次优化，吐出优化后的Graph IR，最后对于不同的后端这些Graph IR都会被编译为特定后端可以识别的机器码完成模型推理。比如对于CPU，NNVM就吐出LLVM可以识别的IR，再通过LLVM编译器编译为机器码到CPU上执行。
融合操作算法基本上是考察当前节点到其支配点所有路径上的节点是否都符合融合规则，融合规则形如Complex-out+Injective/Injective+Injective/Injective+Reduction等。如果符合规则就进行融合，不符合就不融合。
参考文章
TVM系列 - 图优化 - 算子融合_tvm算子融合-CSDN博客
算子融合技术演进与Apollo方案浅析 - 知乎 (zhihu.com)
Op Fusion（一）： 什么是算子融合 - 知乎 (zhihu.com)