Pet: Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections

PET：利用部分等效变换和自动校正优化张量程序

作者/机构: Haojie Wang, Jidong Zhai, Mingyu Gao, Zixuan Ma, Shizhi Tang, Liyan Zheng (清华大学); Yuanzhi Li (卡内基梅隆大学); Kaiyuan Rong, Yuanyong Chen (清华大学); Zhihao Jia (卡内基梅隆大学和Facebook)

A1 主要贡献

核心问题: 现有深度神经网络（DNN）框架通过应用完全等效的程序变换来优化张量程序，即变换前后的子程序在数学上对于任意输入都是等价的。这种方法要求输出张量的每个元素都保持等效性，因此错失了许多潜在的优化机会，因为那些仅在输出张量子集上保持等效性的变换（部分等效变换）被排除了。

研究目标: 本文旨在探索一种全新的张量程序优化方法，即利用部分等效变换来发现现有方法错过的、性能更优的张量程序。为了在应用这些变换的同时保证模型的统计行为不变，需要解决两大挑战：1) 如何快速检查等效性、识别不一致的输出区域并高效生成校正核函数以恢复功能等效性；2) 如何在因引入部分等效变换而急剧增大的设计空间中，有效管理计算复杂性，并结合完全等效变换来发现高性能的张量程序。

创新点:
1. 首次提出利用部分等效变换进行张量程序优化: 本文首次尝试利用部分等效变换并结合自动校正来优化张量程序，探索了比现有DNN框架大得多的搜索空间。
2. 建立严格的理论基础: 提出了一系列严谨的理论基础，极大地简化了等效性检查和校正核函数的生成过程。这使得即使应用了部分等效变换，也能在实践中保持模型的统计行为。
3. 设计高效的生成与优化方法: 提出了一种高效的生成和优化方法，能够自动探索包含完全和部分等效变换的巨大设计空间。
4. 实现端到端框架PET: 将上述技术实现为一个端到端的框架PET，并在五个真实世界的DNN模型上进行了评估。实验表明，PET通过发掘以往被忽略的部分等效变换机会，性能最高可达现有系统的2.5倍。

A3 背景知识与设计原则

现有深度神经网络（DNN）框架（如TensorFlow 【3, Martín Abadi et al. Tensorflow: A system for largescale machine learning. OSDI, 2016.】、TensorRT 【32, NVIDIA TensorRT: Programmable inference accelerator. https://developer.nvidia.com/tensorrt, 2017.】 和 TVM 【6, Tianqi Chen et al. TVM: end-to-end optimization stack for deep learning. CoRR, 2018.】）中，一种常见的优化形式是完全等效变换，它在提升张量程序性能的同时保持其数学等效性。当前完全等效变换的例子包括算子融合 【2, TensorFlow Graph Transform Tool. https: //http://github.com/tensorflow/tensorflow/tree/ master/tensorflow/tools/graph_transforms, 2018.】、布局变换 【18, Chao Li et al. Optimizing memory efficiency for deep convolutional neural networks on gpus. SC'16, 2016.】 和图替换的自动生成 【15, Zhihao Jia et al. Taso: optimizing deep learning computation with automatic generation of graph substitutions. SOSP, 2019.】。尽管这些变换在提升性能方面是有效的，但它们只探索了程序优化的有限空间。

相比之下，图1展示了一个用于卷积算子的部分等效变换示例。它将两个独立的图像沿宽度维度拼接成一个更大的图像以提高性能。这是因为更大的宽度（通常是现代加速器如GPU上卷积的最内层维度）能提供更多的并行性并改善计算局部性。然而，经过这种变换后的新程序（如图1(b)所示）在输出张量拼接边界的子区域上（如图1(b)中阴影框所示）产生了不同的结果，导致了部分不等效。除了通过改变张量形状和线性化来优化张量程序的例子外，部分等效变换还包括用更优化的、具有相似语义的算子替换效率较低的算子，以及修改张量程序的图结构以启用额外的优化。尽管部分等效变换在性能提升方面展现出巨大潜力，但由于其可能影响模型准确性，在当前的DNN框架中并未被考虑。手动实现这类变换是不可行的，因为它需要评估大量潜在的变换以发现有前途的变换，并且为了保持模型准确性，需要校正核来修复不等效部分的结果（见图1(c)）。因此，需要更自动化的方法来发现性能优化的部分等效变换并校正其结果，这也是本文的主要焦点。

图1：一个通过操纵张量形状和线性化来提高卷积性能的部分等效变换。 (b)中的阴影框突出了两个程序之间不等效的元素。(c)中的校正核应用于这些元素以恢复输入程序的功能等效性。
图1：一个通过操纵张量形状和线性化来提高卷积性能的部分等效变换。 (b)中的阴影框突出了两个程序之间不等效的元素。(c)中的校正核应用于这些元素以恢复输入程序的功能等效性。

多线性张量程序（MLTPs）的定义: PET是首个利用部分等效变换并自动校正其结果来优化张量程序的框架，其核心是利用了张量程序的多线性特性。首先定义多线性张量算子：一个具有n个输入张量$I_1, ..., I_n$的算子op是多线性的，如果op对所有输入$I_k$都是线性的。DNN计算通常由多线性张量算子（如矩阵乘法、卷积）和逐元素的非线性算子（如ReLU 【23, Vinod Nair and Geoffrey E. Hinton. Rectified linear units improve restricted boltzmann machines. ICML'10, 2010.】 和 sigmoid）组成，其中线性算子因其高计算复杂性消耗了大部分计算时间。如果一个程序P中的所有算子都是多线性的，则该程序是一个多线性张量程序（MLTP）。

$$op(l_1, \dots, l_{k-1}, X, \dots, l_n) + op(l_1, \dots, l_{k-1}, Y, \dots, l_n) \\ = op(l_1, \dots, l_{k-1}, X + Y, \dots, l_n) \\ \alpha \cdot op(l_1, \dots, l_{k-1}, X, \dots, l_n) = op(l_1, \dots, l_{k-1}, \alpha \cdot X, \dots, l_n)$$

PET框架概览: 如图2所示，PET的输入是一个待优化的张量程序。PET首先将输入程序分割成较小的子程序，以减少每个子程序的探索空间。对于每个子程序，PET的变异生成器（mutation generator）通过为子程序中的MLTPs生成可能的变异体（mutants）来发现部分等效变换，每个变异体具有与原始MLTPs相同的输入和输出形状。为了保持与输入程序的端到端等效性，PET的变异校正器（mutation corrector）会检查变异体与原始MLTP之间的等效性，并自动生成校正核来修复变异体的输出。最后，校正后的变异体被送入PET的程序优化器（program optimizer），该优化器将现有的完全等效变换与部分等效变换相结合，构建一个全面的程序优化搜索空间，并评估每个子程序的丰富变异体集，应用跨边界的后优化，以发现搜索空间中高度优化的候选项。

图2：PET概览。
图2：PET概览。

A2 方法细节

4. 变异生成器

本节描述PET中的变异生成器，它以一个MLTP为输入，并自动生成可能替代输入MLTP的变异体。该生成算法能够发现达到特定大小的有效变异体。每个生成的变异体不一定在整个输出张量上与输入程序保持数学等效性。为了恢复功能等效性，变异校正器（§5）会自动生成校正核。

4.1 变异生成算法

变异体定义与生成: 如果MLTP $P_1$ 和另一个MLTP $P_0$ 具有相同数量的输入（和输出），并且每个输入（和输出）具有相同的形状，我们称 $P_1$ 是 $P_0$ 的一个变异体（mutant）。$P_0$ 和 $P_1$ 的计算不一定等效。直观地说，如果 $P_0$ 是一个张量程序中的子程序，那么用 $P_1$ 替换 $P_0$ 会产生一个有效但可能不等效的张量程序。对于一个给定的MLTP $P_0$，PET使用一组给定的多线性算子O作为基本构建块来生成 $P_0$ 的潜在变异体。表1列出了我们评估中使用的算子，涵盖了各种常用的张量算子，包括计算密集型算子（conv, matmul等）、逐元素算子（add, mul等）和张量操作算子（split, transpose等）。这个集合也可以扩展以包含新的DNN算子。

表1：PET中使用的多线性张量算子。
表1：PET中使用的多线性张量算子。

深度优先搜索算法: 算法1展示了一个用于构建MLTP $P_0$ 潜在变异体的深度优先搜索算法。PET从一个没有算子、只有 $P_0$ 原始输入张量集合的空程序开始。PET通过枚举来自O的算子类型和算子的输入张量，迭代地向当前程序P添加新算子。输入张量可以是 $P_0$ 的初始输入张量（即算法1中的$I_0$）或先前算子的输出张量。深度优先搜索算法枚举所有达到特定大小（称为变异深度）的潜在MLTP。对于每个变异体P，PET检查P和$P_0$是否具有相同数量和形状的输入/输出。如果通过此测试，则P是一个有效的变异体。

算法1：MLTP变异生成算法。
算法1：MLTP变异生成算法。
4.2 示例变异体类别

自动发现的变异类别: 虽然上述变异生成算法足够通用，可以探索足够大的设计空间，但我们强调几类对PET特别重要并能带来性能提升的变异体。需要注意的是，PET不依赖于手动指定的类别，而是自动发现这些类别。

Reshape和Transpose: 张量的内存布局在优化张量程序中扮演着重要角色，这一点广为人知【6, Tianqi Chen et al. TVM: end-to-end optimization stack for deep learning. CoRR, 2018.】。PET利用reshape和transpose算子来变换输入张量的形状和张量维度的线性化顺序，以生成性能更好的变异体。reshape算子通过将单个维度解耦为多个维度或将多个维度组合为一个来改变张量的形状。transpose算子修改张量维度的线性化顺序。reshape和transpose通常联合应用于变换张量布局。例如，图1展示了一个卷积算子的潜在变异体，它将两个独立的图像沿宽度维度连接起来（即图1(b)中的$T_1 \rightarrow T_3$）以提高卷积性能。这种连接涉及三个reshape和transpose算子的组合：首先，一个reshape算子将$T_0$的批处理维度分割成一个内维度（将每两个连续图像分组）和一个外维度；然后，一个transpose算子将新创建的内维度移动到宽度维度旁边，并相应地更新张量的线性化顺序；最后，另一个reshape算子将两个图像组合起来。变异生成器通常将多个连续的reshape和transpose算子融合成一个复合算子，即reshape & transpose，这减少了生成变异体的大小，并允许探索更大、更复杂的变异体。

单算子变异体: PET也可以生成变异体，用一个不同且性能更好的算子替换张量程序中的低效算子。一些标准的张量算子，如卷积和矩阵乘法，在现代硬件后端上已经通过手动或自动方式进行了广泛优化。相比之下，它们的变体，如跨步或扩张卷积【20, Yuhong Li et al. Csrnet: Dilated convolutional neural networks for understanding the highly congested scenes. CVPR, 2018.】，则没有得到同样高效的支持。将它们变异为其具有高度优化核函数的标准对应物有性能上的好处。例如，图3展示了一个将扩张卷积转换为常规卷积的变异体，它通过根据给定的扩张因子重新组织输入张量的线性化顺序。然而，该变异体与输入程序并非完全等效，需要后续校正以恢复功能等效性。

多算子变异体: PET还支持用另一组更高效的算子替换由多个算子组成的子图。例如，几个具有相似张量形状的独立卷积可以组合成一个更大的卷积，以提高GPU利用率并减少核函数启动开销。这需要操作张量形状并添加适当的填充（参见§8.3.3中的示例）。

5. 变异校正器

校正器的目标与挑战: PET生成的变异体虽然性能可能更高，但可能会在输出张量的某些区域产生不同的数学结果，从而可能导致精度损失。为了在应用层面上保持透明性，PET选择保持输入程序的统计行为并保证相同的模型精度，这得益于变异校正器。具体来说，变异校正器以一个MLTP $P_0$ 及其一个变异体P为输入，并自动生成校正核，这些核函数应用于P的输出张量以保持与 $P_0$ 的功能等效性。变异校正器的目标有两个：首先，对于任何给定的MLTP及其变异体，校正器分析这两个程序，并识别输出张量中所有结果相同的区域，这些区域因此不需要任何校正。其次，对于其余输出不同的区域，校正器自动生成核函数来修复变异体的输出并保持功能等效性。设计变异校正器需要解决两个挑战：一是输出张量可能非常大，涉及多达数百万个需要等效性验证的元素，逐一验证是不可行的；二是每个输出元素的验证可能依赖于许多张量算子中的大量输入变量，数值枚举所有可能的输入值是不切实际的。

核心理论: 两个极大地简化验证任务的定理是PET变异校正器的核心。PET无需针对所有输入值组合来验证所有输出位置，而只需验证少数代表性输出位置和少量随机生成的输入值。这极大地减少了验证工作量。

5.1 理论基础

MLTP输出的数学表示: 为简化分析，我们假设输入的MLTP $P_0$ 及其变异体P各有一个输出。我们的结果可以通过顺序分析每个输出来推广到多输出程序。让$P(I)$表示在n个输入张量$I = (I_1, ..., I_n)$上运行P的输出张量。让$P(I)[\vec{v}]$表示在位置$\vec{v}$的输出值，让$I_j[\vec{u}]$表示$I_j$在位置$\vec{u}$的输入值。根据这些定义，一个MLTP P的单个输出位置的计算表示为公式，其中$R(\vec{v})$是计算$P(I)[\vec{v}]$时的求和区间，而$\vec{u} = L_j(\vec{v}, \vec{r})$是一个从$(\vec{v}, \vec{r})$到第j个输入张量$I_j$的位置$\vec{u}$的线性映射。例如，一个核大小为3×3且零填充的卷积定义为公式(1)，其中D、H和W分别指输入图像$I_1$的通道数、高度和宽度。两个线性映射可以表示为$L_1(\vec{v}, \vec{r}) = (d, h+x, w+y)$ 和 $L_2(\vec{v}, \vec{r}) = (d, c, x, y)$，其中$\vec{v} = (c, h, w)$ 和 $\vec{r} = (d, x, y)$。

图片 P(I_1,...,I_n)[v] = Σ(r∈R(v)) Π(j=1 to n) I_j[L_j(v,r)]
图3：一个将扩张卷积转换为标准卷积的示例变异体。(b)中的红色阴影框突出了两个程序之间的不等效元素，这些元素由(c)中的校正核修复。
图片 conv(I_1, I_2)[c,h,w] = ...

Box的定义: 输出张量的不同位置可能有不同的求和区间。对于上面定义的卷积算子，计算左上角输出位置（即h=0, w=0）只涉及一个2×2的核（即$0 \le x \le 1, 0 \le y \le 1$），因为该位置没有左边或上边的邻居，如图4所示。我们将具有相同求和区间的输出位置分组为一个box。形式上，一个box是输出张量的一个区域，其中所有元素都具有相同的求和区间。这个卷积总共有九个box，如图4所示。同一个box中的所有输出位置具有相同的求和区间并共享相似的数学性质，PET在检查程序等效性时利用了这一点。PET无需在所有单个位置上测试两个MLTP的等效性，而只需在每个box中验证m+1个特定位置，其中m是输出张量的维度数。

图4：一个3×3核、零填充的卷积的九个box及其求和区间。一个卷积有三个求和维度（即方程(1)中的d, x, y）。通道维度(d)在所有box中具有相同的区间，因此被省略。
图4：一个3×3核、零填充的卷积的九个box及其求和区间。一个卷积有三个求和维度（即方程(1)中的d, x, y）。通道维度(d)在所有box中具有相同的区间，因此被省略。

定理1: 对于两个具有m维输出张量的MLTP $P_1$ 和 $P_2$，让$\vec{e}_1, ..., \vec{e}_m$是一组m维基向量。即$\vec{e}_i = (0, ..., 0, 1, 0, ..., 0)$是一个m元组，除了第i个坐标为1外，所有坐标都为0。设B是$P_1$和$P_2$的一个box，设$\vec{v}_0$是B中的任意位置。定义$\vec{v}_j = \vec{v}_0 + \vec{e}_j, 1 \le j \le m$。如果对于所有输入I和所有$0 \le i \le m$，$P_1(I)[\vec{v}_i] = P_2(I)[\vec{v}_i]$成立，那么对于所有输入I和所有$\vec{v} \in B$，$P_1(I)[\vec{v}] = P_2(I)[\vec{v}]$也成立。证明概要: 证明使用了一个引理，即如果$P_1$和$P_2$在位置$\vec{v}_0$和$\vec{v}_0+\vec{e}_i$上是等效的，那么等效性对$\vec{v}_0+k \cdot \vec{e}_i$也成立，其中k是一个整数。我们通过比较$P_1$和$P_2$相对于输入变量的系数矩阵来证明这个引理。利用这个引理，我们证明$P_1$和$P_2$在整个box B上是等效的，因为任何$\vec{v} \in B$都可以分解为$\vec{v}_0$和$\vec{e}_0, ..., \vec{e}_m$的线性组合。定理1表明，如果$P_1$和$P_2$在一个box中的m+1个特定位置上是等效的，那么等效性对同一box中的所有其他位置也成立。这极大地减少了验证工作量：PET只需在每个box中验证m+1个特定位置，而不是检查输出张量的所有位置。

定理2: 尽管如此，单个位置的验证仍然具有挑战性，因为每个MLTP通常涉及大量输入变量。证明两个MLTP的等效性需要检查这些输入变量所有可能的值分配组合。我们使用以下定理进一步解决这个挑战：对于两个具有n个输入张量的MLTP $P_1$和$P_2$，设$\vec{v}$是一个$P_1$和$P_2$不等效的位置，即存在I使得$P_1(I)[\vec{v}] \ne P_2(I)[\vec{v}]$。设$I_0$是一个从有限域F中均匀采样的随机输入。$P_1(I_0)[\vec{v}] = P_2(I_0)[\vec{v}]$的概率最多为$n/p$，其中p是F中可能值的数量。证明概要: 这是Schwartz-Zippel引理【28, Jacob T Schwartz. Fast probabilistic algorithms for verification of polynomial identities. JACM, 1980.】【35, Richard Zippel. Probabilistic algorithms for sparse polynomials. International symposium on symbolic and algebraic manipulation, Springer, 1979.】的一个推论。定理2表明，如果两个具有n个输入的MLTP在特定位置$\vec{v}$上不等效，那么它们在从有限域F中采样的随机输入上产生相同结果的概率很低。这个定理显示了随机测试在检查两个MLTP等效性方面的充分性和有效性。由于定理2依赖于从有限域F采样随机输入，而MLTP在实数无限域上操作，我们选择F为一个模p的整数域，其中p是一个大素数（在我们的评估中$p = 2^{31} - 1$）。随机测试中的算术运算在整数上进行，并对素数p取模。在有限域上工作还提供了另一个理想的特性，即应用算术运算符不会涉及整数溢出。通过结合定理1和2，PET将原始的验证任务（即在所有输入值组合下检查所有输出位置）简化为一个更轻量的任务，只需使用几个随机生成的输入测试少数代表性位置，如表2所示。

表2：PET中验证工作量的减少。
表2：PET中验证工作量的减少。
5.2 变异校正算法

三步校正算法: PET变异校正算法利用§5.1中的定理来计算变异体输出张量中哪些区域与输入MLTP不等效，因此需要额外校正。具体来说，只需使用少量随机测试来检查来自两个MLTP的每对重叠box的等效性。整个算法分三步进行：
1. 第一步：Box传播: 首先，我们通过box传播计算给定MLTP的box。box传播的思想类似于深度学习中的前向和后向传播：我们根据输入的box计算算子输出张量的box，计算遵循程序中的算子依赖关系。我们为张量的每个维度维护一组分裂点，以识别其box的边界。对于一个多线性算子，我们根据其输入张量的分裂点以及算子类型和超参数来推断其输出张量的分裂点。图5展示了图1中变异示例的box传播过程。
2. 第二步：对每个box对进行随机测试: 在获得输入MLTP $P_1$ 及其变异体 $P_2$ 的所有box后，PET利用§5.1中的定理来检查来自$P_1$和$P_2$的每对box的相交区域。如果两个box没有任何重叠区域，则可以跳过。对于每个box交集，PET在由定理1确定的m+1个位置上检查两个程序的等效性，其中m是输出张量的维度数。对于这m+1个位置中的每一个，PET通过分配从有限域F（包含0到p-1之间的所有整数，$p = 2^{31} - 1$是一个素数）中均匀采样的值的输入张量来进行一组随机测试。因此，两个不等效的MLTP在随机输入上产生相同输出的概率最多为n/p，其中n是MLTP的输入数。最终，两个不等效的MLTP通过所有测试的概率低于$(n/p)^t$，其中t是测试用例的数量，也是PET中的一个超参数，用于在校正器速度和不等效MLTP通过所有随机测试的错误概率之间进行权衡。我们的方法引入了一个极小且可控的错误概率，即不等效的程序可能以$(n/p)^t$的概率通过随机测试。我们认为这是一个随机测试如何在程序验证成本和验证张量程序变换的微小不健全概率之间进行权衡的例子。为了进一步减少验证工作量，PET包含了一个缓存优化：所有box的测试共享同一组随机输入，并且PET缓存并重用所有中间结果以避免冗余计算。
3. 第三步：校正核生成: 对于每个未通过随机测试的box，PET生成校正核来修复其输出，并恢复原始MLTP与其变异体之间的数学等效性。为了修复输出，校正核执行与原始MLTP相同的操作集，但仅在两个输入程序不等效的那些box上（如图1中红色阴影框所示）。这些box在多维空间中是规则的立方体，可以看作是原始张量的子张量，但尺寸要小得多。因此，PET直接利用现有的DNN库【8, Sharan Chetlur et al. cudnn: Efficient primitives for deep learning. CoRR, 2014.】【10, Dense Linear Algebra on GPUs. https://developer. http://nvidia.com/cublas, 2016.】或核生成技术【6, Tianqi Chen et al. TVM: end-to-end optimization stack for deep learning. CoRR, 2018.】【34, Lianmin Zheng et al. Ansor: generating high-performance tensor programs for deep learning. OSDI 20, 2020.】来生成校正核。为了减少校正开销，PET会机会性地将校正核与现有的张量算子融合（§5.3）。

图5：图1示例的Box传播过程。红色箭头表示每个张量维度的分裂点。
图5：图1示例的Box传播过程。红色箭头表示每个张量维度的分裂点。
5.3 校正核的融合

融合以降低开销: 校正核可能会引入不可忽视的开销，原因在于启动校正核的成本及其有限的并行度。例如，一些校正核的执行时间可能与相应的全尺寸张量算子相当，这可能会抵消应用部分等效变换带来的性能增益。为了减少校正开销，PET会机会性地将校正核与其他张量算子融合。例如，图6(b)展示了应用图1中的部分等效变换后的张量程序。Conv-2是用于修复Conv-1输出的校正核。由于这两个卷积算子共享相同的权重（即W1），PET将它们融合成一个单一的卷积，如图6(c)中的Conv-1-2所示。这种融合需要将T1和T'0连接成一个单一的张量，并将Conv-1-2的输出分割成T2和T'2。连接和分割仅涉及直接的内存复制，并且可以与reshape和transpose算子融合。

图6：将校正核与DNN核融合。
图6：将校正核与DNN核融合。

6. 程序优化器

优化器概述: 本节描述PET中的程序优化器，它探索一个结合了完全和部分等效变换的大型程序优化搜索空间，并快速发现高度优化的程序。程序优化器首先将输入程序分割成多个较小尺寸的子程序，以实现高效的变异生成（§6.1）。其次，为了优化每个独立的子程序，PET通过改变要一起变异的算子子集和变异的迭代轮数，在丰富的候选空间中搜索最佳变异体（§6.2）。最后，在将优化后的子程序拼接在一起时，PET在子程序的边界上应用额外的后优化，包括冗余消除和算子融合（§6.3）。整个程序优化算法总结在算法2中。

算法2：程序优化算法。
算法2：程序优化算法。
6.1 程序分割

分割策略与原因: 变异生成的复杂性随输入程序大小迅速增长，因此直接变异一个包含数百个算子的大型张量程序几乎是不可能的。PET转而将输入程序分割成多个不相交的、尺寸更小的子程序。正确选择输入程序的分割点至关重要，既要有效降低变异复杂性，又要保留大部分程序优化机会。我们使用张量程序中的非线性算子作为分割点，原因有三：首先，非线性算子（如DNN中的激活层）在张量程序中广泛使用，通常每一或几个线性算子后会跟一个非线性激活，这有效地将分割后的子程序限制在我们期望的小尺寸内。其次，PET的变异只适用于MLTPs，任何非线性算子都必须从变异中排除，这使得在非线性算子处分割成为一个自然选择。第三，我们的设计也受到一个观察的启发，即大多数现有的张量程序变换【2, TensorFlow Graph Transform Tool. https: //http://github.com/tensorflow/tensorflow/tree/ master/tensorflow/tools/graph_transforms, 2018.】【6, Tianqi Chen et al. TVM: end-to-end optimization stack for deep learning. CoRR, 2018.】【15, Zhihao Jia et al. Taso: optimizing deep learning computation with automatic generation of graph substitutions. SOSP, 2019.】在其替换模式中也不包括非线性算子（除了融合，我们在§6.3中处理）。在非线性算子处分割后，PET会进一步调整子程序大小。对于多个没有数据依赖的独立子程序，PET会考虑使用分组或批处理算子将它们组合成一个单一子程序。另一方面，如果一个子程序仍然过大，PET每次只会用其算子的一个子集查询变异生成器（见§6.2）。

6.2 子程序优化

搜索最佳变异体: 将输入程序分割成多个独立的子程序后，PET通过查询§4.1中的变异生成器来变异每个子程序，并将估计性能最好的前K个候选项保存在一个堆结构H中，如算法2的第7至16行所示。在每一步，获得的每个变异体都会替换当前每个候选项（即算法2中的P）中对应的子程序，以生成一个新的候选项（即$P_{new}$），然后对其应用一系列后优化（见§6.3）。PET使用一个改编自TASO【15, Zhihao Jia et al. Taso: optimizing deep learning computation with automatic generation of graph substitutions. SOSP, 2019.】的成本模型来估计每个新候选项$P_{new}$的性能。该成本模型对每个配置的张量算子测量一次执行时间，并通过累加其算子的测量执行时间来估计新程序候选项$P_{new}$的性能。

管理变异过程: 为了在合理的时间和空间成本内探索每个子程序足够大的可能变异体空间，我们通过几个关键特性来管理变异过程。首先，当一个子程序中的算子数量超过阈值d（我们的评估使用d=4）时，PET通过枚举所有最多包含d个算子的可能组合，将子程序分解成更小的算子子集，并仅在该子集上查询变异生成器，同时保持其余算子不变（算法2第26行）。其次，我们允许对一个子程序进行最多r轮的迭代变异（算法2第23行），这显著扩大了可能变异体的搜索空间，并使PET能够发现更优化的变异体。所有轮次中生成的所有变异体都作为潜在候选项返回给优化器。值得注意的是，PET的优化器与现有的完全等效变换【2, TensorFlow Graph Transform Tool. https: //http://github.com/tensorflow/tensorflow/tree/ master/tensorflow/tools/graph_transforms, 2018.】【15, Zhihao Jia et al. Taso: optimizing deep learning computation with automatic generation of graph substitutions. SOSP, 2019.】兼容，并且可以将其纳入。通过结合完全和部分等效变换，PET探索了一个显著更大的程序优化搜索空间，并发现了现有优化器错过的深度优化程序。

6.3 后优化

跨边界优化: 最后，所有子程序的优化变异体需要被拼接在一起。除了连接它们的输入和输出张量，我们还在子程序边界之间执行几种后优化，以进一步提高整体性能。我们观察到PET中的变异生成器会引入大量的reshape和transpose（R/T）算子，特别是在每个子程序的开始和结束处。这为跨子程序融合这些R/T算子，并进一步融合从上述子程序优化中排除的非线性算子提供了机会。如图7所示，为了优化子程序之间的边界，PET首先通过重新排序R/T算子与逐元素的非线性激活（如ReLU和sigmoid），将子程序之间的所有R/T算子分组在一起（图7(b)）。这种重排在功能上是正确的，因为reshape和transpose都与逐元素算子是可交换的。重排还允许PET将非线性激活与其他线性算子融合，例如将一个Conv和一个后续的ReLU融合成一个Conv-ReLU（图7(c)）。然后我们应用以下三种后优化：

  • 逆操作消除 (Inverses elimination): 我们消除任何可以相互抵消的R/T算子对，它们等效于无操作。我们称每一对这样的算子为一个逆组（inverse group），并直接移除它们作为后优化的一部分。图7(b)中的R/T-E和R/T-G就是一个逆组的例子。
  • 算子融合 (Operator fusion): 如图7(c)所示，PET将剩余的连续R/T算子融合成一个单一的算子（例如R/T-DH），以减少核函数启动成本。张量程序中的非线性激活也会与R/T或其他线性算子融合。
  • 预处理 (Preprocessing): 如果一个算子的所有输入张量都是静态已知的，我们就对其进行预处理。例如，在图7(b)中，R/T-B和R/T-I都可以在卷积权重张量w1和w2上进行预处理。
图7：拼接两个子程序SG-1和SG-2时应用的后优化。R/T指一个reshape后跟一个transpose。Conv和ReLU分别表示卷积和ReLU算子。
图7：拼接两个子程序SG-1和SG-2时应用的后优化。R/T指一个reshape后跟一个transpose。Conv和ReLU分别表示卷积和ReLU算子。

7. 实现

端到端框架: PET被实现为一个端到端的张量程序优化框架，包含约13,000行C++代码和1,000行Python代码。本节描述我们对PET变异生成器和校正器的实现。

变异算子: 表1列出了PET当前实现中包含的张量算子。我们使用cuDNN 【8, Sharan Chetlur et al. cudnn: Efficient primitives for deep learning. CoRR, 2014.】和cuBLAS 【10, Dense Linear Algebra on GPUs. https://developer. http://nvidia.com/cublas, 2016.】作为我们的后端算子库。PET也可以扩展以包含其他库，如TVM 【6, Tianqi Chen et al. TVM: end-to-end optimization stack for deep learning. CoRR, 2018.】和Ansor 【34, Lianmin Zheng et al. Ansor: generating high-performance tensor programs for deep learning. OSDI 20, 2020.】。在我们的评估中，我们展示了在TVM和Ansor上的这种可扩展性，并表明它们可以直接受益于PET的部分等效优化和自动校正。Reshape和transpose是部分等效变换中两种常用的算子。我们的实现包括对它们的一系列优化，包括消除R/T算子的逆组和融合连续的R/T算子，如§6.3所述。由于reshape和transpose都是多线性算子，PET直接使用§5中介绍的随机测试方法来检查一系列R/T算子是否形成一个逆组，从而可以被消除。

校正核: §5.2描述了一种通用的方法，通过在结果不正确的位置上直接运行原始程序来生成校正核。为了减少校正开销，PET将校正核与其他张量算子融合，如§5.3所述。校正核融合引入了额外的内存复制，这些复制在后优化过程中也与R/T算子融合。

A4 实验环境

  • 硬件配置:
    • CPU: 双路28核 Intel Xeon E5-2680 v4 处理器（启用超线程）。
    • 内存: 256 GB DRAM。
    • GPU: 1块 NVIDIA Tesla V100 GPU。
  • 软件配置:
    • CUDA/cuDNN: CUDA 10.2 和 cuDNN 7.6.5。对于TVM和Ansor的实验，直接使用这些后端生成的最佳核函数。
    • 输入格式: PET接受ONNX模型作为输入。对于TensorFlow和TensorFlow-XLA，使用onnx-tensorflow工具进行格式转换。
  • 模型架构:
    • Resnet-18: 广泛用于图像分类的卷积网络【14, Kaiming He et al. Deep residual learning for image recognition. CVPR, 2016.】。
    • CSRNet: 用于语义分割的扩张卷积网络【20, Yuhong Li et al. Csrnet: Dilated convolutional neural networks for understanding the highly congested scenes. CVPR, 2018.】。
    • Inception-v3: GoogLeNet的改进版，用于提升准确率和计算效率【31, Christian Szegedy et al. Rethinking the inception architecture for computer vision. CVPR, 2016.】。
    • BERT: 用于自然语言任务的语言表示架构【12, Jacob Devlin et al. BERT: pre-training of deep bidirectional transformers for language understanding. CoRR, 2018.】。
    • Resnet3D-18: 用于视频处理的3D卷积网络【13, Kensho Hara et al. Learning spatio-temporal features with 3d residual networks for action recognition. ICCV Workshops, 2017.】。
  • 实验参数:
    • 默认变异生成深度为4（算法1中的depth=4）。
    • 默认搜索轮数为4（算法2中的r=4）。
    • 使用CUDA事件测量端到端推理延迟。

A4 实验结果

端到端性能评估:
* 实验内容: 将PET与现有张量程序优化器（TensorFlow、TensorFlow XLA、TensorRT、TASO）在五种DNN模型上进行端到端推理性能比较，批处理大小分别为1和16。所有优化器均使用相同的cuDNN和cuBLAS后端库，以确保性能差异仅源于优化后的张量程序。
* 实验结果: 如图8所示，PET在所有模型上均取得了最佳性能。即使是对于Resnet-18这类已被高度优化的模型，PET仍能带来高达1.21倍的性能提升。对于较新的模型如CSRNet和BERT，PET的速度最高可达现有框架的2.5倍。对于Resnet-18、CSRNet和Inception-v3，批处理大小为16时，PET的加速比更高，因为更大的批处理大小为PET提供了更多跨张量维度进行变异的机会。
* 分析结论: PET通过发现现有优化器未曾考虑的新的部分等效变换，显著提升了DNN模型的推理性能。

图8：PET与现有框架的端到端性能比较。对于每个DNN，PET柱上方的数字显示了相对于最佳基线的加速比。TASO不支持Resnet3D-18中的3D卷积算子。
图8：PET与现有框架的端到端性能比较。对于每个DNN，PET柱上方的数字显示了相对于最佳基线的加速比。TASO不支持Resnet3D-18中的3D卷积算子。

部分等效变换的有效性验证:
* 实验内容: 将PET发现的部分等效变换及其对应的校正核手动添加到TASO中，作为额外的图替换规则，并测量增强版TASO的性能提升。
* 实验结果: 如图9所示，增强版的TASO在Inception-v3和BERT上的推理性能分别提升了1.12倍和1.31倍。
* 分析结论: 这表明部分等效变换确实扩大了图变换的设计空间，而PET能够自动发掘这些优势。PET通过校正核融合和后优化，能够避免一些因校正开销过大而无法被TASO利用的变换。

图9：将PET的部分等效变换添加到TASO后的性能收益。
图9：将PET的部分等效变换添加到TASO后的性能收益。

案例研究:
* 张量级优化: 在Inception-v3的一个卷积算子上，PET通过变换输入张量形状，使得更高效的FFT卷积算法得以应用，从而将运行时间减少了7倍，GPU DRAM和L2访问分别减少了100倍和15倍（见表4）。
* 算子级优化: 在CSRNet中，PET将一个实现效率较低的扩张卷积替换为常规卷积，从而启用了更高效的Winograd算法，执行时间减少了1.94倍（见表4）。
* 图级优化: PET为Inception-v3发现了新的图变换，如图10所示。例如，通过填充权重张量，将两个并行的卷积融合成一个组卷积，从而提高计算效率。PET还能发现现有框架错过的完全等效变换。
* 核函数融合: 在CSRNet的优化中，如果不进行校正核融合和R/T算子融合，最终程序的性能会下降2.9倍，甚至比原始程序更慢。这证明了PET中融合优化的关键性（见图11）。

表3：算子基准测试列表。
表4：关于表3中conv和dilatedconv算子性能的案例研究。
图10：PET为Inception-v3发现的变异体。axis表示执行concat和split的维度。
图11：PET对CSRNet的优化细节。(a) 优化前的CSRNet。(b) 优化后的CSRNet。(c) DilatedConv-4在子程序优化后、后优化前的子程序。(d) 未融合的R/T算子及F R/T-4对应张量的形状。

与TVM和Ansor的结合:
* 实验内容: 评估PET与TVM和Ansor等核生成技术的结合效果。在cuDNN/cuBLAS、TVM、Ansor三种后端上，分别测试有无PET优化的一组常用DNN算子的性能。
* 实验结果: 如图12所示，将PET与TVM和Ansor结合，相比于直接为这些算子生成核函数，性能可分别提升高达1.23倍和1.21倍。
* 分析结论: PET的程序级优化与现有的核生成技术是正交的，可以结合使用以获得更好的性能。

图12：PET在cuDNN/cuBLAS、TVM和Ansor后端上的性能比较。性能相对于未使用PET的cuDNN/cuBLAS进行了归一化。
图12：PET在cuDNN/cuBLAS、TVM和Ansor后端上的性能比较。性能相对于未使用PET的cuDNN/cuBLAS进行了归一化。

消融和敏感性研究:
* 实验内容: 评估仅使用完全等效变换、仅使用部分等效变换以及两者结合（即PET）的性能。同时，研究变异深度和变异轮数这两个超参数对优化结果性能的影响。
* 实验结果: 如图13所示，仅使用等效变换的性能增益与TASO等先前工作相似；仅使用部分等效变换能带来显著好处但仍有潜力未被发掘；而PET通过结合两者获得了最高性能。如图14所示，增加变异深度和轮数通常能带来性能提升，但收益会逐渐饱和。例如，将深度从2增加到3对两个模型都有显著提升，因为PET发现的许多变异体都是包含三个算子的子程序。
* 分析结论: 联合考虑完全和部分等效变换是PET取得最佳性能的关键。PET的搜索复杂度适中，但能获得显著的性能增益。

图13：使用仅（完全）等效变换、仅部分等效变换以及两者（如PET中）的张量程序优化性能比较。
图14：使用不同变异深度（§4.1）和轮数（§6.2）的PET性能比较。

搜索时间:
* 结果: PET发现高度优化的程序变异体通常需要不到3分钟（Resnet-18为89秒，CSRNet为88秒，BERT为91秒，Resnet3D-18为165秒）。优化Inception-v3耗时约25分钟，因其模块分支较多。
* 结论: PET的搜索时间与TASO和Ansor等先进的DNN优化框架相当，作为部署前的一次性成本是可以接受的。

A5 结论

本文提出了PET，这是第一个利用部分等效变换和自动校正来优化张量程序的DNN框架。PET能够发现那些仅具有部分功能等效性但能提升DNN计算性能的程序变换，并借助严谨的理论保障，后续通过自动校正恢复完全等效性。评估结果显示，PET通过解锁现有框架错过的部分等效变换机会，性能最高可达现有框架的2.5倍。PET已在 https://github.com/thu-pacman/PET 公开源。

A6 附录

A.2 Pet用法

API与后端支持: Pet提供了C++ API来构建输入张量程序，并且支持从ONNX模型导入。对于每个输入张量程序，Pet会生成一个数学上等效的可执行文件，其中包含了本文描述的性能优化。Pet默认使用cuDNN/cuBLAS作为后端，但用户也可以导出变异子程序及其对应的输入/输出张量形状，以使用不同的后端，如TVM和Ansor。

A.3 范围

可复现性: 该构件可用于评估和复现论文的主要结果，包括端到端评估、算子级评估、不同优化策略和启发式参数下的性能比较，以及搜索时间。

A.4 内容

实验内容: 构件评估包括以下实验：
* E1: Pet与其他框架的端到端性能比较。(图8)
* E2: 在不同后端（cuDNN/cuBLAS、TVM、Ansor）上的算子级性能比较。(图12)
* E3: 不同优化策略（完全等效变换、部分等效变换、两者联合优化）的性能比较。(图13)
* E4: 使用不同启发式参数的性能比较。(图14)
* E5: 搜索时间。(第8.6节)

A.5 托管

代码库: 此构件的源代码可在GitHub上找到：https://github.com/whjthu/pet-osdi21-ae，master分支，提交ID为9e07cb1 。

A.6 需求

硬件依赖: 此构件依赖于NVIDIA V100 GPU。
软件依赖: 此构件依赖以下软件库：
* Pet使用cuDNN和cuBLAS库作为后端。我们的评估使用CUDA 10.2和cuDNN 7.6.5。
* 在E1和E2中，TensorFlow、TensorRT、TASO、TVM和Ansor被用作基准DNN框架。我们对这些基准的评估使用了TensorFlow 1.15、TensorRT 7.0.0.11、提交ID为f11782c的TASO（我们为TASO添加了一些小修复以支持测试模型），以及提交ID为3950639的TVM。

A.7 安装

从源码安装Pet: 克隆git代码，创建build目录，执行cmake ..和make -j，然后设置PET_HOME环境变量。
安装其他框架: 请参考构件评估说明（git仓库中的README.pdf）或框架提供的安装说明。

A.8 实验工作流

实验细节: 此构件包含了以下实验。所有DNN基准测试均使用GPU设备内存中的合成输入数据，以消除CPU和GPU之间数据传输的副作用。详细的运行说明可在构件评估说明中找到。
* A.8.1 端到端性能 (E1): 此实验复现论文中的图8。需要先生成ONNX模型，然后分别在各自的目录下运行TensorFlow & XLA、TensorRT、TASO和PET的测试脚本。
* A.8.2 算子级性能 (E2): 此实验复现论文中的图12。脚本位于operator_ae文件夹中。TVM和Ansor的实验会花费很长时间搜索不同的变异核。
* A.8.3 不同优化策略 (E3): 此实验复现论文中的图13。脚本位于pet-ae文件夹中。
* A.8.4 不同启发式参数 (E4): 此实验复现论文中的图14。脚本位于pet-ae文件夹中。
* A.8.5 搜索时间 (E5): 此实验复现论文中的第8.6节。E1中Pet的命令会同时打印搜索时间。注意，由于AE的评估平台CPU与论文中使用的平台不同，搜索时间可能会有差异，但应在同一数量级。