【深度学习编译器系列】4. 前端优化：算子融合

最新推荐文章于 2024-05-16 14:54:01 发布

Elaine_Bao

最新推荐文章于 2024-05-16 14:54:01 发布

阅读量3.1k

点赞数 6

分类专栏：深度学习框架文章标签：深度学习编译器

本文链接：https://blog.csdn.net/elaine_bao/article/details/129506409

版权

本文探讨了深度学习编译器中算子融合的重要性，以提升性能和减少内存读写时间。文章介绍了TVM、XLA、Pytorch、OpenVINO、ONNX和AITemplate等编译器的算子融合策略，包括融合规则和实现方法，并对不同编译器的特点进行了比较总结。

摘要由CSDN通过智能技术生成

在接下来的几篇博客中，我们一起来了解一下前端优化中常见的方法。今天先来看一下算子融合。

1. 为什么要做算子融合

算子融合的核心思想是将多个算子合并为一个，因而无需将中间结果写回全局内存，减少了中间变量的分配，从而提升了性能。另外合并以后只需调用一个kernel，也能减少多个kernel调用的时间。

举个例子，常见的算子融合比如Convolution+ReLU, 如果不做算子融合，Convolution的计算结果需要写回到CPU或GPU的内存里，然后ReLU再从内存里读出来进行计算。由于ReLU本身的计算量挺小的，所以这时候性能的瓶颈就在内存的读写上了。如果做了算子融合，那么ReLU在每次计算完Convolution以后直接inplace就做掉了，这样就减少了中间变量的读写时间，从而提升了性能。

2. 不同深度学习编译器中的算子融合

说完为什么要做算子融合以后，再来说说不同的深度学习编译器都是怎么做算子融合的。

2.1 TVM

在TVM中，算子分为四类：单映射（injective）、规约（reduction）、融合复合式（ complex-out-fusible）和不透明式（opeque）。定义算子时就要确定对应的类别。针对不同的类别，TVM的融合规则如下图所示：

opeque算子不能与其他算子进行融合
complex-out-fusible算子可以和elemwise injective的算子进行融合
injective算子可以和其他injective算子进行融合
reduction算子可以和injective算子进行融合

确定完算子融合的规则以后，TVM的算子融合算法是基于支配树实现的。

算子融合的一个难点在于如何处理潜在的钻石形状分支（diamond shape branches），示例如下图所示。假设conv2d可以被fuse到elemwise add上。但是在conv2d那个点的时候，我们并不知道它的所有输出consumers最后都会汇聚到elemwise add这一个节点。

/*
            conv2d
            /  |  \
           /   |   \
         op    op   op
          \    |    /
           \   |   /
          elemwise add
               |

TVM如何解决这个问题呢？就是利用支配树和支配点的信息。

支配树：由各个点的支配点构成的树。
支配点：是当前节点的所有consumers汇聚的最邻近的节点。比如在上述例子中，elemwise add是conv2d的支配点。

首先，TVM会根据DAG形式的计算图构造出支配树。然后根据支配树信息，进行算子融合。具体来说，遍历图上的每个节点，检查该节点到其支配点之间的所有路径是否符合融合规则，如果符合就对其进行融合。注意，这里的支配点可能已经和其它节点融合过了，此时融合算法依然能够正确运行。

2.2 XLA

XLA的算子融合规则有四类：

Instruction Fusion: 这是一种简单的竖向结构的算子融合。producer instruction会被融合到它的consumer上。XLA维护了一张表，用来指示哪些算子被融合是比较划算的，哪些算子则不是很划算。
Fusion Merger: 这种融合是把producer instruction