TASO

最新推荐文章于 2024-05-18 09:59:17 发布

weixin_37958272

最新推荐文章于 2024-05-18 09:59:17 发布

阅读量1.5k

点赞数

分类专栏：深度学习优化算法

本文链接：https://blog.csdn.net/weixin_37958272/article/details/113870332

版权

深度学习优化算法专栏收录该内容

12 篇文章 2 订阅

订阅专栏

TASO: Optimizing Deep Learning Computation with Automatic Generation of Graph Substitutions

现有的深度神经网络（deep neural network，DNN）框架采用人工设计的图变换来优化DNN的计算图。这种方法忽略了可能的图优化，并且很难扩展，因为新的DNN算子是定期引入的。

我们提出了TASO，第一个自动生成图替换的DNN计算图优化器。TASO将一个算子规范列表(a list of operator specifications)作为输入，并使用给定的算子作为基本构件生成候选替换。所有生成的替换都会通过一个automated theorem prover对算子规范进行正式验证。为了优化一个给定的DNN计算图，TASO执行一个基于成本(cost-base)的回溯搜索，应用替换来找到一个优化的图，这个图可以直接被现有的DNN框架使用。

我们对五个DNN体系结构的评估表明，TASO比现有的DNN框架高出2.8倍，同时所需的人力要少得多。例如，TensorFlow目前包含大约53000行手动优化规则，而TASO所需的operator specifications只有1400行代码。

1 Introduction

深度神经网络（DNN）框架将神经结构表示为一个计算图，其中每个节点是一个数学张量算子（例如，矩阵乘法、卷积等）。为了提高计算图的运行时性能，最常见的优化形式是图替换，将与特定模式相匹配的子图替换为功能等同、性能提高的子图。

现有的DNN框架通过应用由领域专家手动设计的图替换来优化计算图，如图1a所示。例如，TensorFlow、PyTorch、TensorRT和TVM使用贪婪的基于规则的优化策略，并直接在输入计算图上执行所有适用的替换（即规则）[6, 8, 31, 36]. MetaFlow[21]允许替换，可以增加或降低性能，以使等效计算图的搜索空间更大，并使用回溯搜索来探索这个空间，但它仍然需要手动指定的替换。虽然人工设计的替换可以提高DNN计算的性能，但在几个方面还存在不足。

Maintainability. 人工设计图替换需要大量的工程工作。例如，TensorFlow R1.14包括155个替换，这些代码在大约53K行的C++代码中实现。由于不断引入新的算子，维护问题变得更加严重；例如，最近的工作针对不同的图像分类任务提出了depth-wise[19]、分组[38]和转置卷积[16]。TensorFlow r1.14目前包括17个图替换（用4K行代码编写）以优化普通卷积（例如，将其与不同类型的算子融合）。对于现有的方法，支持每一个新的卷积变量都需要类似的实现工作，因为每一个都有稍微不同的语义，并且不能使用现有的替换直接优化。

Data layout. 张量数据可以以各种布局存储在内存中，这种选择对运行时性能有很大影响。最佳布局取决于算子(operator)和硬件。例如，在P100 gpu上，卷积在行主布局（即，最里面的维度是连续存储的）上表现最好，而矩阵乘法在列主布局（即，最外面的维度是连续存储的）上表现最好。在支持4×4矩阵运算的Tesla V100 GPU上，最佳性能可能需要将张量分为4×4块。然而，将布局转换与图替换一起考虑会增加另一个层次的复杂性。例如，如果图替换与特定的布局转换相结合，则它只能提高性能（参见第7.5节）。当前的框架通过将数据布局和图替换视为单独的优化问题来避免这种复杂性，并按顺序解决它们[8，27]，如图1a所示，但是这种分离错过了许多可能的优化机会。

Correctness. Hand-written graph替换容易出错，图替换中的bug会导致计算图的错误[2，4]。同样的问题也出现在编译器优化中，不正确的优化会导致不正确的程序。在编译器文献中，已经投入了大量的精力来正式验证优化( formally verifying optimizations)[7，11，13，23，29，30，33，35]。然而，据我们所知，这些技术还没有被应用到由DNN框架执行的图替换优化中。

在这里插入图片描述

1.1 Our Approach

本文介绍了第一个自动生成图替换的DNN计算图优化器TASO（Tensor algebrage SuperOptimizer）。图1b显示了TASO的概述，它在三个方面不同于现有的框架。首先，TASO只需要算子定义和规范(specifications)，并自动生成图替换，减少了人工设计。其次，TASO使用形式化验证(formal verification)来确保生成的图替换的正确性。最后，TASO联合优化了图形替换和数据布局，显著提高了运行时性能。

Generating substitutions. TASO的图替换生成器枚举给定的DNN算子集（例如cuDNN核[10]）上的所有可能的计算图，直到一个固定的大小，并在一组随机输入张量上执行它们。对随机输入具有相同结果的任意一对计算图被视为候选替换。为了有效地找到所有这样的对，TASO构造了一个哈希表，其中基于随机输入的输出的哈希来存储计算图。

Formal verification. TASO的图替换验证器依赖于用户提供的算子属性来确保生成的图替换的正确性。算子属性利用算子的数学属性，例如卷积的线性。我们使用的43个算子属性的完整列表显示在表2中。我们的计算结果表明，每个算子的一小组性质就足以证明替代的正确性。

从形式上看，我们使用基于一阶逻辑的symbolic representation来对张量算子进行建模，与底层(underlying)张量的大小不可知，并且可以简洁地表达算子属性。验证器使用指定的属性来检查所有生成的图替换的正确性，使用自动化的定理推理器。

我们还提出了一种开发算子属性的方法，它在两个方面帮助开发人员：（1）由图替换生成器引导所需属性的发现；（2）使用小尺寸张量的symbolic execution进一步验证算子属性。在开发过程中，我们发现我们的验证方法在算子规范(operator specification)和图替换生成器的实现中发现了几个错误。

Joint optimization. TASO通过将图替换和数据布局转换集成到一个公共表示(by integrating them into a common representation)中来联合优化它们。TASO使用了MetaFlow的基于成本(cost-base)的回溯搜索算法[21]，并扩展了它的成本模型，以捕获不同数据布局引起的性能差异。在搜索过程中，TASO通过硬件上特定的数据布局来衡量DNN算子的性能。这些单独的度量用于预测具有特定数据布局的整个计算图的性能。

Evaluation. 我们在五个DNN架构上评估TASO。对于由ResNet-50[18]等现有框架优化的广泛使用的dnn，TASO通过使用1400行长的算子定义和规范(operator definitions and specifications 1,400 lines long)，将这些框架的性能与手写规则相匹配。

对于诸如ResNeXt-50[38]、Nas-RNN[39]、NasNet-A[40]和BERT[15]等新的DNN体系结构，TASO通过自动发现新的图替换来优化这些体系结构，比最先进的框架快2.8倍。与顺序优化图替换和数据布局相比，联合优化可以进一步提高性能1.2倍。在所有的实验中，TASO在不到十分钟的时间内发现了一个优化的图，这使得在大规模部署前优化DNN架构时使用它是可行的。

2 Graph Substitution Generator

本节介绍TASO替换生成器，该生成器在给定基本算子列表的情况下自动生成潜在替换。生成算法找到所有有效的替换到一定的大小。

为了找到所有可能的替换，一个简单的方法是测试所有图对的等价性，这需要在图之间进行二次测试。我们采用了编译器超优化[7]的思想，为每个图计算一个指纹，它是一些特定输入上的图输出的散列。如果两个图具有不同的指纹，那么它们肯定是不等价的，因此通过仅比较具有相同指纹的图，TASO显著减少了等价测试的次数。在实验中，我们观察到所有具有相同指纹的图都被TASO验证为等价的。

在这里插入图片描述

2.1 Graph Substitution Definition

图替换由三部分组成：（1）与计算图中的子图匹配的source graph；（2）定义一个功能等同的新子图以取代匹配的子图的target graph；（3）source graph和target graph中输入/输出张量之间的映射关系。图2a显示了一个使用矩阵乘法的关联性的图替换示例。图2b使用concatenation和沿行维度split将两个矩阵乘法融合为一个矩阵乘法。A、 B、C、X和Y标识 source graph和target graph中输入和输出张量之间的映射。

图替换是独立于具体张量形状指定的。例如，图2的替换可以应用于任何具体形状的张量A、B和C。一些算子还依赖于配置参数来确定算子的行为。例如，卷积的参数决定了步长、padding和激活（例如，应用relu函数[28]作为卷积的一部分）；split or concatenation的参数决定了算子应用的轴（determine the axis along which to apply the operator）。

Concatenation and split operators. concatenation和split算子通常用于融合具有共享输入的算子，如图2b所示。split算子将张量沿其参数确定的维度划分为两个不相交的子张量。这是一个复杂的问题，因为split point不能从输入张量或参数推断出来。为了解决这个问题，我们观察到split算子总是在先前的concatenation点上划分张量来“undo”最近的concatenation 算子。我们利用这个事实为split算子定义一个合适的语义。

在这里插入图片描述

图3。一种用共享输入融合矩阵乘法的图替换方法。target graph有一个concat和一个split算子，这两个算子都是沿着矩阵的行维度执行的。每个张量的行维度的拆分树(split tree)显示在灰色框中。

形式上，我们为张量的每个维度维护一个split tree来跟踪concatenation历史。图3显示了图2b中所有张量行维度的split tree。split tree允许替换来恢复分裂点(split point)，而不引入任何额外的参数。我们的方法还支持通过concatenation和split算子的嵌套进行多路concatenation和split。

2.2 Generation Algorithm

对于给定的一组算子规范(operator specifications)，TASO分两步生成潜在的图替换，如算法1所示。

在这里插入图片描述

第一步：枚举潜在的graph并收集它们的指纹 TASO首先通过使用给定的算子集来枚举一定大小的所有潜在图。为了构造一个图，TASO通过枚举算子的类型和算子的输入张量，在当前图中迭代地添加一个算子。输入张量可以是图的初始输入张量（如图2中的A、B和C），也可以是之前算子的输出张量（如图2中matmul和concat算子的输出）。

算法1（第7-18行）显示了一种深度优先搜索算法，用于构造不包含重复计算的所有非循环计算图。我们假设一个图包含重复的计算，如果它有两个算子对相同的输入张量执行相同的计算。生成器会忽略这样的图，因为它们不代表有用的计算图。

对于每个图，我们收集它的指纹，指纹是通过对输入张量求值得到的输出张量的散列。TASO使用随机初始化的张量和一些常量作为输入，以允许查找涉及常量张量的替换，例如单位矩阵（参见第7.3节中的示例）。为了避免在计算指纹时出现浮点错误，所有的张量都用整数表示，遵循[37]中介绍的方法。

由于一个图可以有任意数量的输出张量，哈希函数必须确保指纹与输出张量的任何排列无关。为了保证这个属性，TASO使用两步哈希函数来计算指纹，如下所示。

在这里插入图片描述

其中 $t_i$ 是图 $\mathcal G$ 的输出张量。 $hash_1$ 计算输出张量的状态和内容，包括张量的大小、形状和内容。 $hash_2$ 是一个对称的哈希函数，应用于一组无序的哈希值。

Step 2: Testing graphs with identical fingerprint. 对于具有相同指纹的图，TASO在一组测试用例上进一步检验了它们的等价性。与收集指纹类似，每个测试用例都包含一组随机输入张量，如果两个图为所有测试用例生成等价的输出张量，则两个图通过。与指纹不同，这些测试使用介于−1和1之间的浮点数，如果两个输出张量的输出相差不超过一个小的阈值（在评估中为 $10^{−5}$ ），则将其分类为等效输出张量。对于这个阈值，我们观察到整数测试没有差异。但是，可以使用较小的阈值过滤掉对实数有效但会导致浮点错误的替换。

通过随机测试的每一对图成为候选图替换的source graph和target graph，从测试用例中可以自动推断出 source graph和target graph 中输入输出张量之间的映射关系。然后将所有候选的图替换发送给替换验证器以检查其正确性（第3节），然后修剪以消除冗余替换（第4节）。

到目前为止所描述的算法是通用的，因为它不依赖于特定的张量算子。但是，我们观察到对于DNN，有两个算子需要特殊处理。relu算子28]通常用于DNN应用程序，它为所有负输入返回0。由于relu经常返回0，因此会导致许多多余的替换是有效的。为了防止产生这些替换，generator用一个任意的非线性函数代替relu。

在这里插入图片描述

放大算子通过在张量中填充零来增加张量的大小，这对于融合具有不同核大小的卷积非常有用[21]。然而，零的存在也会导致许多多余的替换。为了克服这一点，生成器只考虑将放大应用于输入张量的计算图，即不考虑另一个算子的输出(not to the output of another operator)。这个限制捕获了explade用于融合卷积的预期用途，同时避免了多余的替换(This restriction captures the intended use of enlarge for fusing convolutions, while avoiding the superfluous substitutions.)。

值得注意的是，以前的工作[7]报告了使用随机测试检查编译器超优化中的代码转换时出现的误报。他们观察到一些不正确的代码转换通过了一组测试用例。在所有的实验中，我们没有观察到任何假阳性案例。我们使用一个测试用例来检验所有具有相同指纹的图对，所有通过测试用例的替换都是正确的，并由替换验证器进行验证。这可能是由于DNN算子的高算术密度（high arithmetic density ）和计算图中缺少分支（if语句）。作为参考，[17]表明，对于只有线性算子的程序，两个不等价程序在随机输入上产生相同输出的概率最多为1/d，其中d是一个变量取值范围（即，在TASO中，d= $2^{32}$ ）。

3 Graph Substitution Verifier

我们正式验证(formally verifying)替换的方法背后的关键思想是使用一阶逻辑中表示的一小组算子属性。这些性质是人工设计和审查(manually written and reviewed)，并通过在小尺寸的张量上符号化地执行算子，确认这些张量尺寸的算子性质得到满足，从而进一步验证。算子属性的开发是在替换生成器发现的替换的指导下进行的（ Development of operator properties is guided by the substitutions discovered by the substitution generator.）。

为了验证的目的，我们使用一阶逻辑来建立张量算子模型，其中算子用其参数和输入张量的函数来表示。例如，conv（s，p，c，x，y）表示应用于张量x和y的卷积算子，其中参数s确定stride，p确定padding mode，c确定activation mode，例如，应用relu激活函数作为卷积算子内核的一部分。例如，没有激活的卷积（由 $A_{none}$ 表示）在其第一个参数中是线性的，这一事实由以下算子属性捕获（其中ewadd表示element-wise张量加法）：

在这里插入图片描述

表2。用于验证的算子属性。算子在表1中定义，属性按它们所涉及的算子分组。逻辑变量w、x、y和z是张量类型，变量a、c、k、p和s是张量类型参数变量a表示concatenation and split的轴( the axis of concatenation and split)，c表示卷积的激活模式，k表示pooling的kernel shape，p表示卷积和pooling的padding mode，s表示卷积和padding的步长。

在这里插入图片描述

表1列出了我们评估中使用的所有算子和张量常量，表2显示了我们评估中用于验证图替换的算子的完整列表。

给定算子的性质，我们使用一阶定理证明器，我们的实现使用Z3[14]-来验证所有生成的替换。这种验证相当于一阶逻辑中的entailment checking，检查算子属性是否包含每个生成的替换的source和target图的函数等价性。

使用一阶逻辑对算子进行建模涉及到一定程度的抽象（例如，不建模张量的形状）。我们发现这种抽象级别适合于验证图替换。我们还注意到，数据布局是为了验证而抽象的，布局不影响算子语义，优化器（第5节）确保一致地使用布局。

开发算子属性的方法。 我们根据需要开发了算子性质，以确定使用迭代过程生成的图替换的正确性。在开发过程中，我们运行了替换生成器并尝试验证所有发现的替换。如果一个替换无法验证并且看起来是正确的，我们添加了一个或多个适当的属性。为了防止算子属性出错，我们使用了进一步的验证步骤。

为了验证算子性质，TASO验证了所有参数值和张量大小的组合的算子性质，在我们的评估中，范围是4×4×4×4。为此，TASO需要Python中每个张量算子的基本符号实现。TASO对小尺寸的张量符号化地执行这种实现，有效地将张量运算细化为符号实数算术表达式，其中激活函数（例如relu）使用未解释的函数建模（where activation functions (e.g., relu) are modeled using uninterpreted functions.）。TASO使用Z3（这里是实数运算理论的SMT求解器）来验证算子属性。用于例如，如果用户试图添加（错误的）属性，说明卷积算子对于所有激活模式（包括relu激活）都是线性的，然后，此检查将显示实际算子不满足此属性。

作为辅助开发过程的additional validation step，TASO检查算子属性集是否一致，并且不包含冗余（即，由其他属性包含的属性），这相当于一阶包含检查。这些检查对于发现错误属性也很有用，并且比小张量尺寸的验证更便宜。

在我们的开发过程中，验证方法揭示了几个微妙的错误。图替换生成器在生成无法验证的替换时发现了一些错误，上述验证步骤揭示了候选算子属性中的几个错误。根据我们的经验，一个新的算子只需少量的工作，通常由一个专家工作几个小时就可以完成。通常，必须为每个算子编写几个属性。在我们的评估中，我们能够使用43个算子属性验证所有743个生成的图替换（见表2）。

4 Pruning Redundant Substitutions

如果图替换被更通用的有效替换所包含，那么它是多余的。本节介绍TASO用于消除冗余图替换的修剪技术。所有的剪枝步骤都保留了所有的优化机会：如果图 $\mathcal G$ 可以使用一系列替换转换为图 $\mathcal G'$ ，那么剪枝后 $\mathcal G$ 总是可以转换为图 $\mathcal G'$ （可能使用不同的转换集）。

在这里插入图片描述

图4.Matmul和Add分别指矩阵乘法和元素加法。对于每个子图，A、B、C指的是它的输入张量，而X指的是输出张量。

Input tensor renaming. TASO消除了与其他替换相同的图替换，并对输入张量进行重命名。例如，图4a通过用A重命名输入张量C显示了与图2a等效的冗余替换。对于通过输入张量重命名而等价的替换，TASO修剪了除了一个最一般的替换之外的所有替换。

Common subgraph. TASO还试图消除source and target graph图具有共同子图的替换。TASO确定了两种可能导致剪枝的公共子图形式。

公共子图的第一种形式如图4b所示。在这里，source and target graph都包含一个具有相同输入张量的公共算子（在灰色框中突出显示）。Common subgraph表示source and target graph中其他算子的输入。因此，我们可以用一个新的输入张量来代替公共子图，从而得到更通用的替换。如果这个更通用的替代确实有效，那么TASO就会对不太通用的替换进行修剪。

公共子图的第二种形式如图4c所示。在这里，公共子图（在灰色框中突出显示）包括source and target graph中的所有输出。在这种情况下，通过完全删除公共子图，使其输入成为source and target graph的新输出，从而得到一个更通用的替换。

图4c展示了第二种形式的共同子图。在这里，共同子图（灰色框中突出显示）包括source and target graph中的所有输出。在这种情况下，可以通过完全删除共同子图，使其输入成为 source and target graph的新输出，从而得到一个更通用的替换。

在这里插入图片描述

表3显示了TASO修剪技术对替换数量的影响。我们观察到，这两种修剪技术在消除冗余替换中都起着重要作用，它们的组合使TASO必须考虑的替换数减少了39倍。

5 Joint Optimizer

我们现在描述用于联合优化数据布局和图替换的TASO优化器。优化器使用MetaFlow[21]基于成本(cost-base)的回溯搜索算法，通过应用验证的替换来搜索优化的计算图。TASO扩展了MetaFlow的搜索算法，在执行替换时也考虑了可能的布局优化机会。

在匹配子图上应用替换时，基于source graph中张量的数据布局和算子支持的布局，TASO枚举target graph中张量的可能布局。结果，对匹配的计算图应用替换可能导致具有相同图结构但不同数据布局的多个图。

在这里插入图片描述

图5。一种利用矩阵乘法的转置进行的图替换。matmul和transpose分别表示矩阵乘法和转置。括号表示source and target graphs中每个张量的潜在布局，其中C和R表示张量的列主布局和行主布局。

例如，图5显示了潜在的计算图，这些图可以通过将矩阵乘法的转置应用到具有默认列主布局（如C所示）的source graph上而得到。矩阵乘法和转置算子也都支持替代的行主布局（如R所示）。target graph（即A、B和X）中所有映射张量的数据布局必须与source graph中的布局匹配。target graph中的两个中间张量可以是行主张量，也可以是列主张量，因此在应用这种替换时，TASO会考虑四个不同的计算图（即两个中间张量的CC、CR、RC和RR）。这允许TASO在执行图替换时捕获潜在的布局转换机会。

在这里插入图片描述

算法2展示了我们的基于成本的回溯搜索算法，用于联合优化替换和数据布局成本模型的动机是DNN算子执行无分支的稠密线性代数，因此，在给定相同的数据布局和配置参数（例如，卷积的stride和padding）的情况下，它们在硬件上的性能是高度一致和可预测的。与MetaFlow[21]类似，TASO对每个配置和数据布局都会测量一次DNN算子的执行时间，并通过对其算子的测量执行时间进行加总来估计图形的性能。

为了寻找一个优化的图，所有候选图都保持在一个优先级队列 $\mathcal P$ 中，并按代价递增的顺序出队列。对于每个出列图(dequeued graph) $\mathcal G$ ，TASO考虑每个已验证的替换和适用于替换的可能布局，并应用它们获得功能等价的新图 $\mathcal G'$ 。

在这里插入图片描述

图6。在计算图中引入一个循环的图替换示例，其中A，B是输入，X是输出。原始图计算A×relu（A×B），新图是应用图2b中所示的替换的结果，该替换使用concatenation and split将两个矩阵乘法融合在一起。The source and target graphs of the substitution显示在虚线框中。原始图和替换都是非循环的。但是，生成的图形包含一个循环（以红色突出显示）。

\r\n

图替换的一个非显而易见的属性，可以将循环引入到一个图形中。图6显示了一个示例，其中应用一个有效的替换将导致循环图。由于计算图必须是非循环的，所以TASO在将 $\mathcal G'$ （算法2的第12行）排队到P中之前检查它的非循环性。

最后通过搜索算法返回最佳发现图 $\mathcal G_{opt}$ 。搜索空间被一个超参数α剪枝，它直接消除了所有代价是最佳发现图α倍的图。参数α在搜索时间和最佳发现图之间进行权衡。设置α=1将搜索减少为一个简单的贪心算法，而没有回溯，α的值越高，搜索将探索更多可能的候选，并导致更多的回溯。我们观察到α=1.05在我们的评估中取得了良好的效果。

6 Implementation

TASO被设计和实现为一个通用的、可扩展的张量计算图优化器，这样就可以很容易地添加新的张量算子。表1列出了当前TASO实现中包含的张量算子。有些算子还依赖于其他参数来确定算子的行为，例如stride、padding和卷积的激活。除了算子，TASO还包括四种常张量，它们在代换中很有用。特别地， $I_{ewmul}$ 、 $I_{matmul}$ 和 $I_{conv}$ 是element-wise乘法、矩阵乘法和卷积的单位张量，分别。 $C_{pool}$ 允许将平均池算子转换为深度卷积（参见第7.3节中的示例）。

如第3节所述，TASO使用用户指定的算子属性来验证生成的图替换表2列出了用于验证我们评估中所有替换的43个属性。

TASO可以很容易地扩展到包含新的张量算子。对于对于每个算子，TASO都需要两种形式的输入：（1）算子的参考实现，以及（2）算子属性的规范（specifications of operator properties）。（1）由替换生成器使用的具体实现（C++）和用于验证算子规范（operator specifications）的符号实现（在Python中）组成。根据我们的经验，添加一个新的算子需要专家花费几个小时的工作。

对于目前缺少specifications的新算子，TASO将其视为不透明算子（opaque operator），并且仍然可以使用验证的替换优化图的其余部分。

TASO是在MetaFlow之上实现的，并且重用了MetaFlow基于成本的回溯搜索[21]。总的来说，我们的TASO实现包含大约8000行核心组件（即替换生成器、验证器和优化器）的代码，以及1400行算子参考实现的代码，包括43个算子属性。

TASO是框架无关的，只需以目标框架的输入格式发出优化的图，就可以插入到现有的DNN框架中，如TensorRT和TVM。在评估中，我们在TensorRT和TVM上演示了这种可移植性，并表明它们可以直接使用TASO的优化来提高性能。