EinNet: Optimizing Tensor Programs with Derivation-Based Transformations

• 文章标题：EINNET：使用基于派生的变换优化张量程序
• 作者/机构：Liyan Zheng, Haojie Wang, Jidong Zhai, Muyan Hu, Zixuan Ma, Tuowei Wang, and Shuhong Huang, Tsinghua University; Xupeng Miao, Carnegie Mellon University; Shizhi Tang and Kezhao Huang, Tsinghua University; Zhihao Jia, Carnegie Mellon University

A1 主要贡献

核心问题：现有的深度神经网络（DNN）张量计算优化方法仅考虑由一组固定的预定义张量算子所能表示的变换，这导致了高度受限的优化空间。无论是算子级别的优化器（如TVM、Ansor）还是图级别的优化器（如TASO、PET），它们都局限于所谓的“预定义算子可表示（Predefined Operator Representable, POR）”变换。这种限制使得许多潜在的、更高效的计算重组方式无法被发现，因为它们可能需要现有库中未定义的定制化算子。

研究目标：本文旨在突破POR变换的限制，探索一个更广阔的通用张量代数变换空间。目标是开发一种能够自动发现和利用这些通用变换的优化器，通过揭示算子的计算语义，将计算重组为任意的张量表达式，并自动生成变换所需的新算子，从而发现现有框架无法找到的新型程序变换。

创新点：
本文提出了EINNET，一个基于派生的张量程序优化器，其核心创新点如下：
1. 扩展优化空间：将优化空间从POR变换扩展到通用张量代数变换。通过将算子计算语义和计算图与张量代数表达式相结合，EINNET能够探索一个远大于现有工作的搜索空间，该空间将POR变换作为其特例包含在内。

2. 基于派生的变换发现机制：提出了一种基于派生的机制来自动发现通用表达式之间的变换机会。EINNET通过应用一系列派生规则，将一个表达式变换为其等价的替代形式，从而能够系统地探索广阔的表达式空间，而无需像超优化方法那样枚举所有可能的图。
3. eOperator（表达式即算子）：引入eOperator来表示那些无法用预定义算子表示的非POR计算。这使得EINNET能够生成并利用高度定制化的算子，从而实现新颖的程序变换，例如：(1) 修改算子的计算语义以提高效率；(2) 用高效算子和定制eOperator替换低效算子；(3) 激进地重组计算图以启用后续优化。
4. 距离引导的搜索算法：为解决通用张量代数变换空间的巨大搜索挑战，EINNET采用了一种两阶段的距离引导搜索方法。该方法结合了探索性派生和收敛性派生，利用“表达式距离”来指导搜索过程，使其能够高效地发现需要长序列派生规则的复杂优化。
5. 系统实现与性能：实现了EINNET系统，它在七个真实世界的DNN模型上进行了评估，并在NVIDIA A100和V100 GPU上，相较于现有的张量程序优化器，分别取得了最高2.72倍和2.68倍的性能提升。

A2 动机与概述

EINNET概览。图2展示了EINNET的概览。对于一个输入的张量程序，EINNET首先将其分割成多个由预定义算子组成的子程序。每个子程序由程序转换器翻译成张量代数表达式（§3）。然后，EINNET的基于派生的优化器使用不同的派生规则，包括表达式内和表达式间的派生规则（§4）以及表达式实例化规则（§5），为每个表达式生成优化的子程序，这些子程序包含预定义算子和eOperator。最后，EINNET为每个子程序选择最佳的变换，并对表达式进行后优化，以构建一个高效的张量程序（§6）。

动机示例。图3(a)展示了一个由EINNET发现的优化案例。它首先执行表达式内派生，将卷积变换为矩阵乘法，然后执行表达式间派生，将多个算子融合成一个。图中红色的算子，如OffsetReduce、DLT（数据布局变换）和OffsetReduce+Relu，是EINNET自动发现和生成的eOperator。图3(b)展示了EINNET对图3(a)中Conv3x3算子进行的新优化细节。图3(c)则展示了经典的im2col【36, Aravind Vasudevan等人，Parallel multi channel convolution using general matrix multiplication, 2017】卷积优化，该优化在现有库中广泛实现，并且也被EINNET的自动优化空间所覆盖。与im2col中复制输入张量（次数同核大小）不同，新发现的变换复制输出张量的次数相同。当输出尺寸小于输入尺寸时，这种方法可能更高效，并且在我们的评估中，对于ResNet-18【19, Kaiming He等人，Deep residual learning for image recognition，CVPR 2016】中的某些卷积，在NVIDIA A100 GPU上相比cuDNN实现了2倍的加速。

现有优化器的局限性。现有的张量程序优化器无法自动发现此类变换，原因在于：(1) 这些变换需要eOperator（例如，带有偏移量地添加中间张量），这超出了基于超优化的框架（如TASO【20, Zhihao Jia等人，Taso: optimizing deep learning computation with automatic generation of graph substitutions，SOSP 2019】和PET【38, Haojie Wang等人，Pet: Optimizing tensor programs with partially equivalent transformations and automated corrections，OSDI 2021】）所探索的POR变换空间；(2) 这些变换修改的是计算语义而非调度，因此无法被基于调度的优化器（如TVM【7, Tianqi Chen等人，TVM: an automated endto-end optimizing compiler for deep learning，OSDI 2018】和Ansor【40, Lianmin Zheng等人，Ansor: generating high-performance tensor programs for deep learning，OSDI 2020】）找到。

A3 方法细节

3. 张量代数表达式

EINNET将张量程序表示为张量代数表达式，它定义了如何从输入张量计算输出张量的每个元素。图4展示了三个矩阵相乘（即A × B × C）的表达式。我们现在描述表达式的组成部分。为简单起见，我们假设一个表达式只有一个输出。EINNET的表达式可以轻松推广到多个输出。

遍历和求和符号。遍历符号，表示为$L_{x=x_0}^{x_1}$，由一个迭代器x和一个迭代空间[$x_0$, $x_1$)组成。遍历符号对应于输出张量的一个维度，其中迭代空间是该维度的范围。遍历符号的顺序表示输出张量的布局。例如，在图4中，$L_{c=0}^{C}$后跟$L_{r=0}^{R}$表示表达式的输出是一个形状为C × R的二维张量。

求和符号。求和符号，表示为$\sum_{x=x_0}^{x_1}$，计算在维度x上迭代求和，其中x的取值为{$x_0, x_0 + 1, ..., x_1 - 1$}，为简洁起见，下文表示为范围[$x_0, x_1$)。需要注意的是，一个EINNET表达式在不同求和符号顺序下被认为是相同的，但对应于表达式的不同调度。因此，这被排除在表达式搜索空间之外。

张量索引。张量通过多个迭代器的算术组合进行索引，包括加(+)、减(-)、乘(*)、除(/)和模(%)。为简单起见，我们可以将多个迭代器合并为一个迭代器向量，其迭代空间可以用一个整数集表示或在表达式中省略。例如，$L_{c=0}^{C} L_{r=0}^{R}$可以表示为$L_{cr}$或$L_{\vec{x}}^X$，其中$\vec{x} = (c, r)$是迭代器向量，X = C × R是迭代空间。

Scope（作用域）。对于包含多个算子（例如，两个连续的矩阵乘法A × B × C）的张量程序，一个常见的优化是实例化并重用中间结果（例如，缓存A × B的输出），以避免对这些结果的重复计算。EINNET引入scope来表示中间结果的实例化以便后续重用。形式上，一个张量代数表达式如果其输出被实例化为一个张量，则它是一个scope，用周围的{}表示，这允许后续计算引用该张量，从而避免重复计算。在图4(b)中，对应于A × B的表达式是一个scope，允许后续计算直接引用该表达式的输出。EINNET的许多派生规则都是基于scope之间的变换，包括从现有scope生成新scope、将一个scope转换为另一种形式以及将多个scope合并为一个（§4）。scope之间的变换对EINNET的优化至关重要。

Padding（填充）。某些计算会访问输入区域之外的数据，我们称之为padding。例如，一个3 × 3的卷积可能会有padding。如果未指定，padding值设为0。

通用格式。我们将一个单scope表达式表示为：
$$\mathbf{L} \sum_{\vec{x}}^{\mathbb{X}} \sum_{\vec{y}}^{\mathbb{Y}} f(\mathbf{T}[\tau(\vec{x}, \vec{y})]))$$
其中 T = {T0, T1, ...} 是一系列输入张量，$\tau(\vec{x},\vec{y})$是索引函数，它使用迭代器$\vec{x}$和$\vec{y}$计算T中张量的元素索引，f是作用于T中被索引元素的计算。

4. 派生规则

为了给输入的张量程序发现高度优化的表达式，EINNET使用派生规则对输入表达式进行变换。表1总结了EINNET使用的派生规则。值得注意的是，派生规则的数学等价性保证了EINNET发现的派生表达式的等价性。与旨在发现特定硬件上给定表达式的优化调度的核生成器调度原语不同，EINNET的派生规则专注于变换张量表达式的计算语义，例如将计算重组为高效的算子。

4.1 表达式内派生

表达式内派生规则的作用。表达式内派生规则将一个表达式转换为其他功能等价的形式，这对于为张量程序构建一个全面的表达式搜索空间至关重要。图5展示了图3(b)中的优化细节。它将Conv3x3的表达式分为两部分，将一部分派生为预定义的Matmul算子，然后将另一部分转换为eOperator。我们现在描述这些表达式内派生规则。

求和分裂（Summation splitting）。该规则将一个求和符号$\sum_{\vec{s}}$分解为两个独立的求和$\sum_{\vec{s1}}$和$\sum_{\vec{s2}}$，并通过将内部求和的结果转换为一个scope来实例化它：
$$ L \sum_{\bar{x}, \bar{s}_{1}, \bar{s}_{2}} f(T[\tau(\bar{x}, \bar{s}_{1}, \bar{s}_{2})]) \mapsto L \sum_{\bar{x}, \bar{s}_{1}} (L \sum_{\bar{s}_{2}} f(T[\tau(\bar{x}, \bar{s}_{1}, \bar{s}_{2})]))[\bar{x}, \bar{s}_{1}] $$
其中τ是从$(\vec{x},\vec{s1},\vec{s2})$到输入位置的映射。EINNET将一个求和的迭代器分为两个不相交的组$\vec{s1}$和$\vec{s2}$，这将求和分裂为两个嵌套的scope S1和S2，其中S1是上述表达式中高亮的部分，而$S2 = L_{\vec{x}} \sum_{\vec{s1}} S1[\vec{x},\vec{s1}]$。在求和分裂中，EINNET将内部求和的结果转换为一个scope，其输出被外部求和重用。为了将一个3×3卷积转换为一批九个矩阵乘法（如图5所示），EINNET首先通过将求和$\sum_{crs}$分裂为两个求和$\sum_{rs}$和$\sum_{c}$，将初始表达式E1转换为E2，并实例化内部求和的输出（即 L_{rshw}^f \sum_c A[h + r, w + s, c]K[r, s, f, c] ）。内部scope仅沿c维度求和；结果，一个中间的五维张量被实例化，因为沿r和s维度的求和未执行而被转换为遍历符号。外部scope计算剩余的r和s维度的求和，产生一个三维张量。图5(a)和(b)显示了计算图的变化。

变量替换（Variable substitution）。该规则通过应用双射函数$\Phi$（即$\vec{y} = \Phi(\vec{x})$）将一组遍历迭代器$L_{\vec{x}}$替换为一组新的迭代器$L_{\vec{y}}$。这种变换允许表达式使用一组不同的遍历迭代器进行计算。具体来说，对于表达式$L_{\vec{x}} f(T[\tau(\vec{x})])$，变量替换引入一个中间scope，该scope计算$L_{\vec{y}}^Y f(T[\tau(\Phi^{-1}(\vec{y}))])$，其中$\Phi$是一个将迭代空间X映射到$\Phi(X)$的双射函数，$\Phi^{-1}$是$\Phi$的反函数：
$$\hat{L}_{\vec{x}}^{x} f(T[\tau(\vec{x})]) \rightarrow \hat{L}_{\vec{x}}^{x} \left\{ \hat{L}_{\vec{y}}^{\Phi(X)} f(T[\tau(\Phi^{-1}(\vec{y}))]) \right\} [\Phi(\vec{x})].$$
变量替换会构造一个带有新遍历迭代器的中间scope。为了保持功能等价性，原始迭代器$\vec{x}$被用来利用中间scope的输出构造最终结果。尽管一个表达式存在大量可能的变量替换，EINNET通过分析表达式中的索引函数并检查它们是否能形成双射来推断合法的替换。在图5中，EINNET应用变量替换，使用一个将(r, s, f, h + r, w + s)映射到(r, s, f, t1, t2)的双射函数$\Phi$，将表达式从E2转换为E3。具体来说，h + r和w + s在E3中被t1和t2替换。为了在所有替代方案中自动识别有前景的变量替换，§6.1引入了表达式距离，一种有效探索搜索空间的新技术。

遍历合并（Traversal merging）。该规则使用索引函数$\Phi$将两个独立scope中的遍历符号合并到一个scope中：
$$\underset{\vec{x}}{\overset{X}{\mathbf{L}}}\underset{\vec{y}}{\overset{Y}{\sum}}\left\{\underset{\vec{z}}{\overset{Z}{\mathbf{L}}}f(T[\tau(\vec{z})])\right\}[\Phi(\vec{x}, \vec{y})] \rightarrow \underset{\vec{x}}{\overset{X}{\mathbf{L}}}\underset{\vec{y}}{\overset{Y}{\sum}}f(T[\tau(\Phi(\vec{x}, \vec{y}))])$$
其中索引函数$\Phi$将外部scope的迭代器$\vec{x},\vec{y}$映射到内部scope的迭代器$\vec{z}$，并满足$\Phi(X \times Y) \subseteq Z$。在图5的例子中，EINNET应用遍历合并将E4转换为E5。对于这个变换，外部遍历和求和符号以及内部遍历符号都包含五个迭代器（即$\vec{x} = (h, w, f)$，$\vec{y} = (r, w)$，和$\vec{z} = (r, s, h, w, f)$）。遍历合并应用了一个单位映射函数$\Phi$和一个索引函数$\tau(r, s, h, w, f) = (r, s, f, h + r, w + s)$。遍历合并移除了一个scope并保持了相同的计算图。

边界放松与收紧（Boundary relaxing and tightening）。边界收紧检查对于任意输入，如果某些输出元素是常量，是否可以避免对这些元素的计算。EINNET在表达式上执行常量传播来处理表达式中的常量数和张量中的填充。如果一个输出区域具有常量值，EINNET将其转换为张量的一个属性以避免不必要的计算。相反，边界放松通过添加额外的填充和冗余计算来扩大张量，以探索更多的优化。图6展示了将一个Conv5×5填充为一个Conv6×6，然后将其转换为一个输出通道数翻四倍的Conv3×3的优化。以下公式展示了放松和收紧的执行方式：

$$\underset{\bar{x}}{ \overset{x}{\mathbf{L}}}f(\mathrm{T}[\tau(\bar{x})]) \iff \underset{\bar{x}}{ \overset{x'}{\mathbf{L}}}f(\mathrm{T}[\tau(\bar{x})]),$$
其中$X \subset X'$，并且T在$X' \setminus X$中具有一个常量值。为了限制该规则可能的候选参数数量，EINNET将边界放松和收紧到一个共同的常量。在图5的运行示例中，E4中的公式对t1和t2执行边界放松，将其范围从[r, H + r)和[s, W + s)分别转换为[-1, H + 1)和[-1, W + 1)，因为对于3 × 3的卷积核，r和s的范围是[-1, 1]。边界放松后，计算图从图5(b)转换为(c)。如果存在多个计划，则选择最严格的一个以防止额外的冗余计算。同时，EINNET仍然能够通过多次应用该规则找到引入更多冗余的变换。EINNET执行边界收紧将E5转换为E6。在E5中，由于在t1 = -1, t1 = H, t2 = -1和t2 = W上执行的计算落入张量A的填充区域，计算结果也为零。因此，t1和t2的范围分别从[-1, H + 1)和[-1, W + 1)收紧到[0, H)和[0, W)。边界收紧后，计算图从图5(c)转换为(d)。

派生搜索空间。派生规则使EINNET能够探索一个巨大的表达式搜索空间。图7展示了一个3x3卷积的派生搜索空间。通过应用不同的派生规则，初始表达式被派生为各种等价表达式，如图7中的计算图所示。图5所示的动机示例由派生路径(a) → (b) → (c) → (d) → (e)确定。该图还显示了EINNET发现的许多其他表达式：通过将(d)中的表达式派生为Conv1x1而不是Matmul，EINNET在(f)中发现了一个新的表达式。通过合并求和迭代器，表达式(i)采用了一个eOperator来为后续的Matmul连接多个带偏移的输入，这代表了传统的Im2col优化【36, Aravind Vasudevan等人，Parallel multi channel convolution using general matrix multiplication, 2017】。表达式(k)显示了一个分组卷积通过复制其输入与原始卷积等价。表达式(n)和(p)显示了另外两个候选表达式，它们都将3 × 3卷积分解为更小的卷积，同时使用派生的eOperator保持输出不变。

4.2 表达式间派生

表达式间派生规则。现在我们介绍EINNET中用于分裂、合并和融合表达式的表达式间派生规则。

表达式分裂（Expression splitting）。该规则通过将原始表达式的遍历空间X划分为两个子空间X1和X2，将一个表达式分为两个独立的表达式，其中X ⊆ X1 ∪ X2：
$$ \underset{\dot{x}}{{\overset{x}{L}}}f(T[\tau(\bar{x})]) \implies \underset{\dot{x}}{{\overset{x_1}{L}}}f(T[\tau(\bar{x})]) \sim \underset{\dot{x}}{{\overset{x_2}{L}}}f(T[\tau(\bar{x})]) $$
其中~表示两个表达式的独立性。

表达式合并（Expression merging）。该规则是表达式分裂的逆过程。它通过合并它们的遍历空间X1 ∪ X2 ⊆ X，将两个具有相同计算的独立表达式合并：
$$ \underset{\bar{x}}{\overset{x_1}{\mathcal{L}}} f(\text{T}[\tau(\bar{x})]) \sim \underset{\bar{x}}{\overset{x_2}{\mathcal{L}}} f(\text{T}[\tau(\bar{x})]) \implies \underset{\bar{x}}{\overset{x}{\mathcal{L}}} f(\text{T}[\tau(\bar{x})]) $$

表达式融合（Expression fusion）。该规则使用以下规则将多个依赖的表达式融合成一个：
$$ \stackrel{Y}{\mathbf{L}}_{g(\underline{y})}(\mathbf{T}'(\pi(\mathbf{\overline{y}}))) \circ \stackrel{X}{\mathbf{L}}_{f(\underline{x})}(\mathbf{T}(\tau(\mathbf{\overline{x}}))) \Rightarrow \stackrel{Y}{\mathbf{L}}_{g(\underline{y})}(\mathbf{L}_{f(\underline{x})}(\mathbf{T}(\tau(\mathbf{\overline{x}}))) | \pi(\mathbf{\overline{y}})) $$
其中T'等于上述表达式中高亮部分的计算结果，$E1 \circ E2$表示表达式E2的结果作为输入提供给表达式E1。

应用示例。图3(a)展示了一个涉及表达式间派生的一系列派生过程。EINNET首先应用表达式内派生规则将Conv3x3和Conv1x1转换为两个Matmul和一个eOperator。由于这两个Matmul共享相同的输入和计算模式，EINNET能够对它们应用表达式合并规则。如虚线框所示，EINNET对两个权重张量进行转置和连接，作为Matmul的输入。Matmul的输出被分割以获得两个等价的输出。此外，EINNET应用表达式融合规则来执行垂直算子融合，这是一种将一连串算子融合成单个核以减少数据移动和核启动开销的优化。在图3(a)的实线框中，EINNET通过应用表达式融合将内存密集型算子（例如OffsetReduce和Relu）融合成一个eOperator。

5. 表达式实例化

尽管EINNET可以将所有表达式视为eOperator，并使用现成的核生成器（例如，我们实现中的TVM）来生成可执行程序，但这样做会导致次优性能。这是因为现有的供应商提供的张量库，如cuDNN【10, Sharan Chetlur等人，cudnn: Efficient primitives for deep learning, 2014】和cuBLAS【11, Dense Linear Algebra on GPUs, 2016】，包含一系列高度优化的张量代数核，其性能优于张量编译器生成的对应物。手动调优和自动生成核之间的性能与表达能力权衡带来了挑战和机遇：我们应该机会性地将一些表达式降级到供应商提供的核以实现其性能优势，并使用核生成器为剩余的表达式生成可执行程序。我们将此任务称为表达式实例化。EINNET考虑两种用于表达式实例化的派生规则：(1) 算子匹配允许EINNET机会性地使用现有的高度优化核（例如，cuDNN和cuBLAS）以实现高性能，以及(2) eOperator生成能够为任意eOperator灵活生成核。应用这些规则后，实例化的scope在原始表达式中被张量替换，并与后续的派生分离。

实例化策略。为了将表达式降级为核，EINNET使用一种策略，将计算密集型表达式映射到预定义算子，并为内存密集型表达式使用核生成器。这种策略使EINNET能够受益于现有的供应商库并保持较低的编译时间，因为内存密集型表达式在现有的代码生成框架【7, Tianqi Chen等人，TVM: an automated endto-end optimizing compiler for deep learning，OSDI 2018】中通常涉及较小的调度空间。虽然更积极地利用核生成器有可能超越机会主义策略，但这会在程序优化期间为数百万个可能的表达式引入显著的核调优开销。这是因为在没有实际调优和分析的情况下，很难准确估计一个核的性能。

强度分析。为了确定一个表达式是计算密集型还是内存密集型，EINNET分析其算术强度，该强度计算为其FLOPs与张量大小的比率。算术强度低于阈值（在我们的评估中为4）的表达式被认为是内存密集型的eOperator。EINNET根据此度量决定是执行算子匹配还是为此表达式生成eOperator。下面介绍这两种实例化规则。

5.1 算子匹配

挑战。将一个表达式映射到一个预定义算子是具有挑战性的，因为一个算子可以用多种表达式来表示。例如，虽然图8(a-b)中的表达式形式不同，但它们都可以被实例化为cuBLAS中的批处理矩阵乘法核，因为它支持具有灵活数据布局的张量。
$$C[b,m,n] = L \sum_{k} A[b,m,k] B[b,k,n]$$
$Z[u,v,w] = L \sum_i X[u,0,v,x+1]Y[u-1,u,x,w]$
$B[b,k,n] = B[l_{b0}b+l_{k1}k+1 \cdot n]$
$Y[u-1,u,x,w] = Y[(l_{y0}+l_{y1})u+l_{y2}x+1 \cdot w - l_{y0}]$

迭代器映射表。EINNET使用一个迭代器映射表来确定一个给定的表达式是否可以映射到一个预定义算子，其中每个算子的迭代器根据其是否出现在算子的输入/输出张量中进行分组。表2显示了几个具有两个输入和一个输出张量的算子的迭代器映射表，包括逐元素算子、批处理矩阵乘法、卷积和G2BMM【23, Johannes Langer, Band Matrices in Recurrent Neural Networks for Long Memory Tasks, 2018】（通用到带状矩阵乘法）。表中的每一行对应一个算子，而每一列显示一个迭代器组。迭代器映射表还记录了每个张量索引中迭代器的系数，用于算子匹配。它可以扩展以支持具有任意数量输入和输出的算子。

匹配流程。迭代器映射表允许EINNET按以下步骤确定一个表达式是否可以映射到一个算子：
1. 匹配张量。为了将给定的表达式映射到一个算子，EINNET枚举了表达式和算子的输入/输出张量之间所有可能的一对一映射。例如，为了将图8(b)中的表达式映射到BMM（即图8(a)中的表达式），存在两种可能的映射，{A → X, B → Y}和{A → Y, B → X}。
2. 匹配迭代器。对于每个张量映射，EINNET进一步使用上述迭代器映射表枚举了表达式和算子之间匹配迭代器的所有可能方式。例如，假设在图8中有一个张量映射{A → X, B → Y}，(b)中的迭代器{u, v, x, w}根据迭代器映射表被映射到(a)中的迭代器{b, m, k, n}（同一组中的迭代器用相同颜色标记）。当同一组中有多个迭代器时，EINNET枚举这些迭代器之间的所有可能映射。
3. 匹配算子属性。许多预定义算子包含用于指定计算的属性。例如，现代BLAS库使用lda和ldb来控制矩阵乘法中输入张量的数据布局。为了匹配这些属性，EINNET将输入和输出张量扁平化（即将它们重塑为一维张量）以隐藏张量形状的复杂性。然后，EINNET通过检查扁平化张量的变量系数来匹配算子属性。图8(c-d)显示了EINNET如何确定BMM算子的属性lB0和lB1。它将两个表达式中的张量B扁平化并比较它们的系数：$l_{B0} = l_{Y0} + l_{Y1}$，$l_{B1} = l_{Y2}$，其中$l_{Yn}$是张量Y的n维度的步幅。表达式(d)中w的系数也被检查以等于表达式(c)中n的系数，因为它们是一对匹配的迭代器。

5.2 eOperator生成

生成流程。对于无法映射到供应商提供的预定义算子的表达式，EINNET将它们转换为eOperator。由于eOperator精确地定义了其计算，EINNET可以直接将其提供给现成的核生成框架（例如，TVM【7, Tianqi Chen等人，TVM: an automated endto-end optimizing compiler for deep learning，OSDI 2018】）。例如，对于图5中的表达式E7，它对应于变换后计算图中的OffsetReduce，EINNET通过将其迭代器空间转换为一个张量，并将计算表达式转换为一个lambda函数，将其提供给TVM。图9显示了EINNET为表达式E7生成的TVM代码，这可以作为TVM的输入程序来生成一个可执行的核。

T1 = tvm.te.compute((H, W, F), lambda h, w, f: tvm.te.sum(T1[r, s, h+r, w+s, f], axis=[r, s]))

图9：由EINNET为表达式E7生成的TVM代码

6. 程序优化器

本节描述EINNET的程序优化器，它使用§4和§5中描述的表达式派生和实例化技术来优化输入的张量程序。这些派生规则创建了一个庞大而复杂的、与输入功能等价的程序搜索空间。EINNET使用距离引导的搜索算法来探索该空间（§6.1），并开发了一种指纹技术来修剪冗余（§6.2）。最后，§6.3描述了EINNET如何协调这些技术来执行端到端优化。

6.1 距离引导的搜索

搜索挑战与策略。为了探索由EINNET派生创建的搜索空间，纯粹的随机搜索策略只能探索有限的路径集或较小的搜索深度，导致次优性能。为了应对这一挑战，EINNET开发了一种两阶段距离引导的搜索算法来应用派生。它包括一个探索性派生阶段和一个收敛性派生阶段，如图10所示。

探索性派生（Explorative derivation）。在探索性派生期间，EINNET迭代地将所有派生规则应用于当前表达式。一个超参数MaxDepth决定了EINNET在探索性派生期间应用的最大派生规则数量。如§5所述，EINNET机会性地使用供应商提供的核库来最大化性能。因此，EINNET利用收敛性派生来快速地将一个表达式派生向一个目标算子（例如，cuDNN和cuBLAS中的算子）。EINNET自动生成必要的eOperator来弥合当前表达式和目标算子之间的差距。

收敛性派生（Converging derivation）。在收敛性派生期间，EINNET首先选择一个目标算子，并使用一种新颖的度量——表达式距离，来指导此阶段派生规则的应用。表达式距离衡量给定表达式E1与给定算子E2的规范表达式之间的差异。为了计算E1和E2之间的距离，EINNET首先使用迭代器映射表（见§5.1）匹配E1和E2中的所有迭代器，并将不匹配的迭代器总数作为它们的距离。具体来说，当前表达式和目标算子之间的每个迭代器不匹配表明两个表达式在某个迭代器组（见表2）中的迭代器数量不同。EINNET应用派生规则来修复不匹配，例如使用变量替换规则来合并/分裂迭代器，从而减小表达式距离。例如，为了将图5中E6内部scope中的表达式派生为Matmul，EINNET比较它们的迭代器（表2）并获得以下匹配：t1, t2 → m; r, s, f → n; c → k。为了修复不匹配，EINNET应用变量替换将迭代器t1和t2合并为m，并将r, s, f合并为n。

自动生成eOperator。在所有迭代器匹配之后，EINNET根据目标算子推断每个输入/输出张量的形状，并通过添加eOperator从现有张量构造新张量。例如，用于Matmul的新输入张量A'和权重张量K'由以下表达式构造：
$A'[m, k] = A'[t_1 \times W + t_2, c] = A[t_1, t_2, c]$
$\mathbf{K}[k, n]=\mathbf{K}^{\prime}[c, r \times S \times F+s \times F+f]=\mathbf{K}[r, s, f, c]$
其中映射函数为$(m, k) = \Phi_A(t1,t2, c) = (t1 \times W + t2, c)$和$(k, n) = \Phi_K(r, s, f, c) = (c, r \times S \times F + s \times F + f)$，W、S和F是迭代器w、s和f的范围。EINNET自动生成表达式(2)和(3)来修复不匹配并减小表达式距离。在收敛性派生期间，EINNET只考虑那些减少当前表达式与目标算子之间表达式距离的派生，这使得EINNET能够修剪大多数派生并快速收敛到目标算子。通过将迭代器映射表中的算子作为目标算子进行枚举，EINNET找到了涉及不同算子的变换。

延迟代码生成。为了加速搜索，EINNET估计派生程序的性能，以避免为eOperator频繁生成代码。具体来说，预定义算子的执行时间通过在硬件上分析其核来测量。同时，eOperator的运行时间根据其输入/输出大小和硬件内存带宽进行估计。我们观察到这种估计是准确的，因为eOperator是内存密集型的，并且通常只占总执行时间的一小部分。

6.2 冗余剪枝

冗余来源与指纹技术。应用不同序列的派生可能会得到相同的表达式。例如，将一个迭代器分裂成两个然后又将它们合并，结果还是原来的那个。为了修剪冗余，EINNET使用一种指纹技术来检测重复的表达式。指纹是表达式的哈希值，可以消除以下冗余来源：
- 求和重排：求和可以重排，例如$\sum_{\vec{x}} \sum_{\vec{y}} f(\vec{x},\vec{y})$等价于$\sum_{\vec{y}} \sum_{\vec{x}} f(\vec{x},\vec{y})$。注意，遍历重排不意味着等价，因为它涉及布局变换。
- 操作数重排：可交换二元运算的操作数可以重排，例如$L_{\vec{x}}(T1[\vec{x}] + T2[\vec{x}])$等于$L_{\vec{x}}(T2[\vec{x}] + T1[\vec{x}])$。操作数重排应同时应用于迭代器计算和张量计算。
- 迭代器重命名：迭代器应通过其迭代空间而非名称来区分，例如$L_{x=0}^N L_{y=0}^M f(x, y)$和$L_{y=0}^N L_{z=0}^M f(y, z)$是等价的，前者中的(x, y)应映射到后者中的(y, z)。
- 张量重命名：由不同scope引入的张量可能具有相同的值。

指纹计算方法。为了消除上述冗余来源，EINNET采用以下方法计算指纹。对于遍历迭代器，EINNET使用其迭代空间及其在当前scope中相对于所有其他遍历符号的顺序作为其指纹。由于考虑了顺序，指纹可以区分具有相同迭代空间但在遍历符号中位置不同的遍历迭代器。对于求和迭代器，EINNET仅使用其迭代空间作为其指纹。因此，在求和重排下的表达式具有相同的指纹。为了解决操作数重排问题，EINNET对可交换操作（如加法）使用操作类型和顺序无关的哈希，对其他操作使用顺序相关的哈希。张量的指纹取决于其来源。对于输入张量，EINNET通过哈希其名称来计算其指纹。对于由scope生成的中间张量，其指纹与产生该张量的表达式的指纹相同。

6.3 端到端工作流

算法 1: EINNET的端到端优化工作流
1: 输入: 一个输入张量程序 P
2: 输出: 一个优化的张量程序 P_opt
3:
4: IR = 表达式间规则集
5: SP = 分割 P 并将子程序翻译成表达式
6: P_opt = ∅
7: for E0 ∈ SP do
8:   Q = {E0}
9:   for E ∈ Q do
10:    for r ∈ IR do
11:      Q = Q + r(E)
12:    Q = Q + DISTANCEGUIDEDSEARCH(E)
13:  将 Q 中的最佳变换添加到 P_opt
14: POSTOPTIMIZATION(P_opt)
15: return P_opt

工作流描述。算法1展示了EINNET以端到端方式优化输入张量程序的工作流。对于一个输入程序，EINNET首先使用非线性激活算子作为分割点将其分割成多个子程序。这是因为激活算子除了融合之外通常不提供进一步的优化机会，这一点已被先前的工作【38, Haojie Wang等人，Pet: Optimizing tensor programs with partially equivalent transformations and automated corrections，OSDI 2021】发现。对于每个子程序，EINNET使用每个算子的规范表达式将其翻译成表达式。由于一个子程序可能包含多个算子，因此也包含多个表达式，EINNET应用表达式间派生规则（第11行），并将每个表达式输入到距离引导的搜索（§6.1）中以执行表达式内派生（第12行）。与将表达式内和表达式间优化集成到一个统一的搜索空间并联合执行不同，这种分离的搜索优先考虑那些能将表达式映射到算子的变换。因此，EINNET能够快速找到有前景的变换，并根据变换后结果的执行时间来修剪不必要的搜索状态。

后优化。最后，EINNET为每个子程序选择发现的最佳表达式，执行后优化，并生成一个优化的张量程序。图11展示了两种后优化类型：eOperator融合和编译时表达式求值。在优化子程序时，EINNET生成eOperator以促进优化变换。在后优化期间，通过应用表达式间派生将连续的eOperator融合成单个eOperator。图11(b)和(c)中的虚线框显示了这种情况。EINNET还检测可在编译时计算的表达式以减少运行时开销。例如，图11中的数据布局变换E2可以在后优化期间处理。

A4 实验

实验环境

• 实现：EINNET由超过23,000行C++和Python代码构建。张量表达式派生系统从零开始实现，并实现了一个张量程序的执行运行时。用户可以直接在EINNET中定义张量程序，也可以加载ONNX格式的现有程序。
• 配置：探索性派生的默认最大搜索深度设置为7。
• 硬件平台：实验在一台配备双Intel Xeon E5-2680 v4 CPU、NVIDIA A100 40GB和V100 32GB PCIe GPU的服务器上进行。
• 软件平台：所有实验均使用CUDA 11.0.2、cuBLAS 11.1.0和cuDNN 8.0.3。
• 工作负载（模型）：评估了七个DNN模型，涵盖了各种机器学习任务：
  • InfoGAN【9】：生成对抗网络，用于学习物体的解耦表示。
  • DCGAN【32】：利用深度卷积结构获取层次化表示。
  • FSRCNN【13】：用于快速图像超分辨率。
  • GCN【30】：图像语义分割模型。
  • ResNet-18【19】：著名的图像分类网络。
  • CSRNet【25】：采用扩张卷积进行拥挤场景分析。
  • LongFormer【6】：改进的Transformer模型，用于处理长序列语言。
    所有模型均采用论文和实现中典型的输入形状，以保证评估的实际意义。

实验结果

端到端性能

与现有框架对比。我们比较了EINNET与当前主流DNN框架的端到端推理性能，包括TensorFlow v2.4、TensorFlow XLA、Nimble、TVM v0.10 with Ansor、TensorRT v8.2和PET v0.1。所有框架使用相同版本的cuBLAS和cuDNN作为后端，并使用FP32数据类型进行计算。对于Longformer中的新注意力算子，我们为TVM、TensorRT、PET和EINNET提供了自动调优的核，并在其他框架中通过einsum实现。图12显示了在NVIDIA A100和V100 GPU上，批量大小为1和16时的结果。

- 性能优势：EINNET在A100上比最佳基准快高达2.72倍，在V100上高达2.68倍。对于CNN（如GCN）和语言模型（如Longformer），EINNET都能将其性能提高2倍以上。即使是已被现有框架高度优化的ResNet-18，EINNET通过图3中展示的新型变换，在V100上仍能实现1.2倍的性能提升。对于CSRNet，EINNET发现了与PET类似的变换，并通过消除额外的转置操作获得了进一步的改进，表明EINNET的派生规则覆盖了PET的优化并能发现更多改进空间。

TF32与Tensor Cores性能。图13展示了在A100上使用TF32数据类型和Tensor Cores时的加速情况。我们创建了一个禁用派生优化的基线EINNET-Base。结果显示，虽然EINNET通常比EINNET-Base和TensorRT带来显著加速，但在ResNet-18等模型上TensorRT表现更好。分析表明，TensorRT除了cuBLAS和cuDNN外还利用了许多高效的GPU核，这是其高性能的一个重要来源，超出了EINNET当前的搜索空间。

优化分析

本节分析EINNET在这些DNN上发现的涉及eOperator的优化，这些优化超出了现有张量程序优化器的范围。

• 变换算子类型：EINNET能将低效算子替换为不同类型的优化算子。在ResNet-18和InfoGAN中，将Conv和ConvTranspose转换为Matmul是有益的。如表3所示，将一个Conv3x3转换为Matmul和一个eOperator (OffsetReduce)，显著减少了GPU DRAM访问（从56.7MB降至10.5MB），并实现了2.7倍的加速。类似地，将一个带步长的ConvTranspose派生为一个Matmul和另一个eOperator，显著减少了L2访问。
• 变换算子属性：EINNET可以利用eOperator变换算子属性。图6展示了将一个5x5卷积的核大小转换为3x3，从而能够使用更适合3x3卷积的先进算法（如Winograd）。为此，添加了一个eOperator来分割Conv3x3的输出并在通道维度上进行带偏移的归约。尽管这引入了额外的计算，但如表3所示，它使得能够使用Winograd算法，进一步减少了计算时间和内存访问。
• 变换张量布局：eOperator允许EINNET通过优化张量布局来加速DNN计算。图14展示了对Longformer的布局优化，它将一个带空洞的G2BMM（通用到带状矩阵乘法）通过重排奇偶行/列，转换为一个非空洞的G2BMM。这大大减少了非连续内存访问，代价是引入两个冗余元素。如表3所示，此变换可将L2缓存访问减少95.9%，执行时间减少78.0%。
  
• 变换图结构：对于图14(d)中的Longformer案例，EINNET通过表达式融合发现中间两个互逆的数据布局变换可以被消除。一个更复杂的例子是GCN模型，它包含重复的空间可分离卷积结构。如图15所示，EINNET首先将卷积转换为Matmul和eOperator，然后将前两个Matmul合并为一个，并将后两个不共享输入的Matmul通过表达式合并和算子匹配规则合并成一个BatchMatmul。这些变换使子图执行时间在A100上（batch size=1）优化了4.9倍。
  

与不同后端的集成

EINNET的表达式空间搜索使其可以与不同的后端（数学库和基于调度的代码生成框架）协作。图16展示了EINNET在cuBLAS/cuDNN、AutoTVM和Ansor后端上的优化效果。对于ResNet-18中的Conv3x3和InfoGAN中的ConvTranspose，将其转换为Matmul和eOperator在所有三个平台上都带来了显著的加速。然而，从Conv5x5到Conv3x3的变换对cuDNN有益，但在AutoTVM和Ansor上即使采用Winograd算法也没有性能提升，这表明为每个后端定制变换是有益的。对于Longformer中的G2BMM，由于其不是当前数学库中的预定义算子，在Ansor上的变换带来了4.54倍的加速。

自动化派生分析

• 搜索深度影响：图17分析了不同最大搜索深度对InfoGAN和Longformer的加速效果。在InfoGAN上，当搜索深度从2增加到4时，性能得到提升。对于Longformer，主要加速来自一个4步派生中发现的变换。结论是，EINNET在适度的深度下即可实现大部分性能提升。
  
• 距离引导派生的有效性：距离引导的派生提供了确定的方向，减少了搜索开销。如图18(a)所示，搜索时间随探索性派生最大深度（MaxDepth）呈指数增长。通过引入收敛性派生，EINNET可以减少探索性派生的深度。如图18(b)所示，在ConvTranspose案例中，原本需要12步探索性派生才能发现目标表达式，而使用收敛性派生后，只需6步探索性派生，大大减少了搜索时间（超过99.0%）。
  
• 表达式指纹的有效性：表达式指纹技术修剪了冗余的搜索状态。图19显示，在派生过程中，指纹技术通过识别和消除重复表达式，有效修剪了98.0%的中间状态，并平均减少了98.2%的搜索时间。
  
• 搜索时间成本：借助距离引导派生和表达式指纹，EINNET对大多数子程序的搜索在两分钟内完成，与TASO和PET等现有框架相当。大多数模型的总搜索时间不超过两小时。由于搜索成本对于一个模型是一次性的，并带来持久的收益，EINNET能够部署在实际生产环境中。

A5 结论

我们提出了EINNET，一个基于派生的张量程序优化器。它将张量程序的优化空间从预定义算子可表示的变换扩展到通用表达式，并能根据变换需求创建新的算子。EINNET在NVIDIA GPU上，其性能最高可达现有最先进框架的2.72倍。

A7 补充细节：相关工作

• 基于规则的方法：如TensorFlow XLA、TensorRT和Grappler，依赖专家手动设计的规则进行优化。它们效率高但需要大量工程努力且只能执行已知的优化。DNNFusion【28】使用基于数学属性的图重写规则，但主要适用于逐元素算子，难以扩展到卷积、矩阵乘法等复杂算子。EINNET通过在表达式层面进行派生，避免了为每个算子手动总结规则的需要。
• 基于超优化的方法：TASO【20】和PET【38】使用超优化技术为一组给定的算子生成图变换。与这些框架相比，EINNET通过张量代数表达式派生，将搜索空间从POR变换扩展到了通用表达式。
• 基于调度的方法：Halide【24】开创了计算与调度分离的思想，被TVM【7】、FlexTensor【42】和Ansor【40】等代码生成框架广泛采用。EINNET与这些基于调度的优化器是正交的，它专注于高层的图变换空间，设计与架构无关的表达式派生规则来重组计算。
• 基于任务的方法：Rammer【27】和Hidet【12】引入了任务（计算和内存操作工作负载的抽象）来优化张量程序。EINNET可以与这些优化器兼容，将它们作为执行和核生成后端。
• 张量化（Tensorization）：TensorIR【14】、AMOS【41】和Glenside【34】旨在为张量加速器生成张量化代码，专注于将计算映射到特定的硬件电路（如NVIDIA TensorCore）。相比之下，EINNET采用表达式匹配来匹配任意的线性张量代数表达式，这更加灵活且模式表达式中可以包含未定参数。

A6 附录

摘要：本附录旨在帮助读者复现论文《EINNET: Optimizing Tensor Programs with Derivation-Based Transformations》中的主要评估结果。

范围：本工件可用于评估和复现论文的主要结果，包括模型级评估、算子级评估，以及关于搜索策略的消融研究和超参数研究。

内容：工件包含以下评估结果：
- E1: EinNet与其他框架的端到端性能比较 (图12)。
- E2: 对§7.3中案例的性能研究 (表3)。
- E3: 在数学库和代码生成框架Ansor上优化前后的算子性能 (图15，应为图16)。
- E4: 不同最大搜索深度下的加速比 (图16，应为图17)。
- E5: 不同MaxDepth下的搜索时间以及有无收敛派生的探索性派生步数 (图17，应为图18)。
- E6: 表达式指纹剪枝的消融研究 (图18，应为图19)。

托管：工件源代码托管于GitHub: https://github.com/zhengly123/OSDI23-EinNet-AE，主分支，提交ID为26a47d9 。

需求：
- 硬件依赖：双路Intel(R) Xeon(R) CPU E5-2680 v4 @ 2.40GHz，NVIDIA A100-PCI-40GB GPU，NVIDIA V100-PCIE-32GB GPU。
- 软件依赖：在Ubuntu 22.04 LTS (Linux kernel 5.15.0-58)上评估。依赖CUDA 11.0.2和cuDNN 8.0.3。需要以下框架作为基准：TensorFlow 2.4, TensorRT 8.0, PET 1.0, Nimble (commit ID bac6d10), TVM v0.10.0。

实验工作流：安装说明和后续实验已包含在本工件中。所有DNN基准测试使用GPU设备内存中的合成输入数据，以消除CPU和GPU之间数据传输的副作用。每个实验（E1-E6）的具体复现指南和脚本均在对应的子目录中提供（例如，OSDI23-EinNet-AE/0_model/用于E1）。