OSDI 2021 PET 论文解读

人生必须的知识就是引人向光明方面的明灯。这篇文章主要讲述OSDI 2021 PET 论文解读相关的知识，希望能为你提供帮助。
今天来阅读一篇OSDI 2021的论文，《PET: Optimizing Tensor Programs with Partially Equivalent Transformations
and Automated Corrections》。

论文链接：https://pacman.cs.tsinghua.edu.cn/~whj/pubs/Pet.pdf
开源代码链接： https://github.com/thu-pacman/PET

【OSDI 2021 PET 论文解读】之前也读过OSDI 2020的《Ansor : Generating High-Performance Tensor Programs for Deep Learning》这篇论文，如果说Ansor是在更微观的角度做代码生成，那么这篇PET就可以说是在更宏观的角度做代码生成了。
无论是 Ansor 还是 PET 我个人认为都相当惊艳，我之前对 Ansor 的论文解读在这个仓库中：https://github.com/BBuf/tvm_mlir_learn 。感兴趣的小伙伴可以读一下，而本篇文章就来阅读一下 PET 。
0x1. 标题和作者
标题可以翻译为：基于部分等价变换和自动校正来优化张量化程序。作者团队来自清华，CMU和FaceBook等。这篇论文的一作王豪杰来自清华大学。后面会介绍到这篇论文在生成突变程序集合时，要维护K个效率最高的突变程序时，使用了ASO中的代价模型和评估方式，所以作者有贾志豪大神也不奇怪。
0x2. 摘要现有的框架在图层做优化一般都是基于等价变换，也就时说变换前后的程序是完全等价的。这里等价的意思是给定相同的输入，那么变换前后的程序一定可以得到相同的输出。而这篇论文挖了一个新坑，即做了一个新的框架PET，在优化过程中允许出现部分等价的变换，并且设计了一个高效的搜索算法去组合完全等价以及部分等价的变换以探索更大的搜索空间。并且最终结果也比较好。
0x3. 介绍这里需要先说明一个名词的含义，统计特性。统计特性的意思就是变换前后程序是完全数学等价的特性。目前TVM，以及TensorFlow，PyTorch，TensorRT等框架的变换优化或者叫 Pass 都是满足这个特性的。而部分等价变换不要求变换前后的程序保持这个统计特性，也就是允许变换后的程序和原程序在相同输入时输出的某些位置的元素是不等的。支持部分等价变换可以（1）改变输入Tensor的shape和排列顺序以提高计算效率（2）使用效率更高的算子代替效率更低的算子（3）对图结构进行变换获得更多的高效优化机会。但要支持部分等价变换也有两个挑战。第一个就是如果直接使用部分等价变换会降低模型精度，因此有必要校正那些不等的张量区域，但是快速识别哪些区域是不等的并且产生校正Kernel是一项很难的任务并且要标注出输出的哪些位置在变换前后是不等的也是一个难题。第二个就是应用了部分等价比变换之后张量程序的搜索空间扩大了，生成候选的张量化程序的算法必须仔细管理其计算复杂度。程序优化器（后面有单独的一节讲）必须平衡部分等价变换带来的好处以及因为它引入的额外开销，并结合完全等价变换来获得高性能的张量化程序。
这篇论文提出了一个全新的部分等价变换来优化张量化程序的框架PET，PET主要由3部分组成。

Mutation generator。突变产生器。对于一个输入张量化程序，这是用来产生部分等价变换的输出张量化程序的。每一个突变程序和输入程序在相同输入的情况下，输出Tensor的形状都是一样的，但某些区域的值可以不一样。
Mutation corrector。突变校正器。PET的突变矫正器检查原始程序和突变程序之前等价性并自动生成校正Kernel。并将校正Kernel应用于输出张量以保证整个变换是符合统计特性的。另外PET也尽可能的融合校正Kernel和张量计算Kernel，以减少因为引入校正Kernel带来的额外开销。检查和校正部分等价变换是非常困难的，因为输出张量可能包含多达百万个元素，并且每个输出元素可能都和大量的输入元素有关。如果挨个去验证，开销会非常大。PET的一个关键贡献就是发现了一套严谨的数学理论，大大简化了这个验证过程（将这个过程的复杂度降低到了常数级别）。不是测试输出张量的所有位置，PET只需要测试几个有代表性的位置就可以完成验证。
**Program Optimizer。**首先一个模型被切分成多个子图，然后对每个子图应用部分等价变换来获得更多的优化机会。最后在整个模型的各个子图的边界部分应用一系列后优化，包括冗余消除和算子融合等以达到整体的最佳性能。

贡献方面其实就是上面这三点，我们先提一下PET在几个模型上进行评估的性能。对于ResNet-18提升了1.2倍，对于CSRNet和BERT提升了2.5倍。
0x4. 背景和想法来源这一节没什么好讲的，感觉和Introduction有一点重复，我们只讲讲图1，帮助大家理解什么是部分等价变换。首先图1长这样：

首先(a)代表一个普通的卷积操作，其中是输入Tensor，它的数据排布可以记作：[b, c, h, w]，即批量大小，输入通道数，输入特征图长度和宽度。然后部分等价变换也就是图(b)通过一个reshape和transpose把图中批量方向的相邻两个特征图拼起来了，也就是这样：[b, c, h, w] -> reshape -> [b / 2, 2, c, h, w] -> transpose -> [b / 2, c , h, w, 2] 。也就是图中的，然后和原卷积核做卷积操作之后得到，再利用reshape和transpose将输出特征图还原为原始的输入特征图大小。注意经过这个变换之后我们发现输出Tensor在边界部分存在和原始卷积的输出Tensor数值不相等的情况，所以还需要对不相等的边界部分进行校正，即?图展示的意思。
后面几节我们会详细讲解如何确定数值不相等的部分是哪些以及如何对这些数值不相等的区域进行校正，现在不理解也没关系。
0x5. 设计总览PET是第一个利用部分等价变换去优化张量化程序的框架，它利用了张量化程序的多线性性质。这里首先要解释一下什么是Multi-linear tensor programs (MLTPs) 即多线性张量化程序，后面我们一律用MLTPs的说法。一个拥有n个输入张量的Op，如果对于每一个输入都是线性的那么就称这个Op是多线性的：

其中X和Y是具有和相同形状的任意张量，是任意标量。深度学习模型一般由线性（Conv，MatMul）和非线性（例如ReLU，Sigmoid）的算子所组成，PET框架中使用的线性算子如Table1所示：

注意这个表是可以扩展的。一个程序是多线性张量化程序（MLTP）当且仅当程序中的所有的Op都是多线性的。接下来我们讲讲PET的设计总览，也就是Figure2。

首先原始的张量化程序输入到PET框架中，然后PET首先把这个程序分解成一些小的子程序来降低每个子程序的搜索复杂度。对于每一个子程序，PET的Mutation Generator会通过为子程序的MLTPs产生可能的变体来发现部分等价变换变体程序。每个变体程序和原子程序拥有相同的输入输出Shape。为了保持端到端的数值正确性，PET的Mutation Corrector检查原始程序和突变程序有哪些区域是不等的，并自动生成校正Kernel进行校正，PET利用了严格的数学理论来简化了这个富有挑战性的任务。
校正后的突变体被发送到 PET 的程序优化器，它将现有的完全等价变换与部分等价变换结合起来，构建一个程序优化的综合搜索空间。优化器为每个子程序评估一组丰富的突变体，并在它们的边界上应用后期优化，以便在搜索空间中发现高度优化的候选程序。
0x6. Mutation Generator这一节主要描述了Mutation Generator的算法实现流程并描述了几种生成的典型突变模式。Mutation Generator的算法如下图所示：

首先有一个原始的多线性张量化程序MLTP