Felix: Optimizing Tensor Programs with Gradient Descent

Felix：使用梯度下降优化张量程序

作者/机构:
Yifan Zhao, Hashim Sharif, Vikram Adve, Sasa Misailovic
University of Illinois Urbana-Champaign, USA

A1 主要贡献

核心问题：在各种硬件上为深度神经网络（DNNs）等张量程序获得高性能实现是一项挑战。现有的基于搜索的张量程序优化器虽然能自动找到高性能程序，但由于搜索空间巨大，其搜索过程效率低下，通常需要数小时甚至数天才能发现好的程序。

研究目标：本文旨在解决现有搜索方法效率低下的问题，提出一种新颖的、基于梯度下降的编译器优化框架，以更高效地找到高性能的程序调度（schedules）。

创新点与主要贡献:
本文提出了Felix，一个新颖的基于梯度的张量程序编译器优化框架。其核心思想是将离散的程序调度搜索问题转化为一个可微的优化问题，从而利用梯度下降方法进行高效搜索。

  1. 可微程序空间的构建：Felix通过对程序空间应用连续松弛（continuous relaxation）技术，并创建了程序延迟的可微估计器，从而构建了一个适合梯度下降搜索的可微张量程序空间。这与传统方法在离散搜索空间上对不可微目标函数进行搜索形成鲜明对比。

  2. 自动生成可微程序特征：Felix能够自动生成与调度变量（如分块因子、循环展开因子）相关的可微程序特征。具体做法是：

    • 创建包含调度变量的符号化调度 (symbolic schedules)。
    • 将符号化调度应用于子图，生成符号化程序 (symbolic programs)。
    • 分析符号化程序，提取出程序特征的数学表达式（公式）。
    • 通过平滑核（smoothing kernels）将这些公式中不可微的操作符（如 min, max）转化为可微的近似函数。
    • 这种方法使得Felix能够应用于广泛的张量算子和编译器变换，因为预训练的DNN成本模型可以一次性训练，然后应用于各种张量程序。

  3. 高效的梯度下降搜索：Felix利用梯度下降来优化调度变量，快速发现高性能的调度方案。梯度下降搜索产生的是连续空间中的候选调度，Felix通过离散化（如四舍五入）并验证其合法性，最终在目标硬件上进行少量实证评估，选出最终的优化调度。

  4. 显著的性能和效率提升：实验结果表明，Felix在平均7分钟的搜索时间内，其优化后的网络性能超过了PyTorch、TensorFlow和TensorRT等现成的推理框架。与最先进的基于搜索的优化器TVM Ansor相比，Felix达到95%最佳性能的速度平均快3.4倍，达到99%最佳性能的速度平均快2.8倍。这表明Felix在时间受限的调优场景或资源受限的边缘设备上尤其有效。

  5. 通用性与实现：Felix在Apache TVM张量编译器及其Ansor优化器之上实现。其优化方法是通用的，可以被其他基于张量的编译器采用。

A3 问题陈述

基于搜索的张量编译器典型工作流程：图1展示了如Ansor、Halide自动调度器和Tiramisu自动调度器等典型基于搜索的张量程序编译器的工作流程。用户首先通过API或声明式语言定义计算（步骤1），编译器将其直接翻译成中间表示，生成一个性能通常较低的初始程序$p_0$。优化器的目标是找到一个程序变换序列$s$（称为调度），生成一个优化后的程序$T(p_0, s)$。
图1. 基于搜索的张量编译器的典型工作流程。
优化问题可以表示为：
$$\max _{s \in S\left(p_0\right)} \operatorname{Performance}\left(\mathcal{T}\left(p_0, s\right)\right)$$
其中，$S(p_0)$是适用于初始程序$p_0$的合法调度搜索空间，$T(p_0, s)$是通过应用调度$s$到$p_0$上生成的优化程序，性能是在目标机器上测量的执行时间。

搜索空间定义：为了获得比手动调优更好的性能，张量编译器需要应用多种优化，如循环分块、循环展开、向量化、并行化以及硬件特定的变换（如GPU上的CUDA块/线程划分）。每种变换都包含可调参数（通常是布尔或整数值），这导致了巨大的搜索空间。为解决这个问题，现有编译器通常使用调度模板或顺序构建并剪枝调度，以限制任意的变换序列。搜索空间依赖于输入程序$p_0$，并且包含离散参数，因为被调优的编译器参数是整数值（如分块大小、SIMD并行因子）。

成本预测模型和程序特征：由于调度的搜索空间是离散的，许多经典的离散空间搜索技术（如遗传算法和束搜索）被用于此优化问题。为了减少对大量候选程序进行实证评估的开销，现有编译器通常使用成本模型（如前馈网络、LSTM、决策树）来预测程序的性能。这些成本模型不直接以调度$s$为输入，而是以从程序$p := T(p_0, s)$中提取的$K$个程序特征向量为输入：
$\text{Feat}(p) = (\text{Feat}_1(p), \ldots, \text{Feat}_k(p), \ldots, \text{Feat}_K(p))$
这些程序特征，例如程序中浮点加/乘操作的数量和循环中缓冲区访问的重用距离，通常由编译器设计者手动选择并通过静态分析提取。

现有目标函数的不可微性：如图1所示，估计一个调度的性能主要有3个步骤（步骤3-5）：生成程序$T(p_0, s)$，提取程序特征$Feat(T(p_0, s))$，并将其输入成本模型。使用这种性能估计器时，优化问题变为（与公式1对比）：
$$\max_{s \in S(p_0)} \text{CostModel} (\textbf{Feat}(\mathcal{T}(p_0, s)))$$
这个目标函数是不可微的，因为传统上应用调度和提取程序特征的过程都是不可微的程序。为了解决这些问题，Felix通过生成符号化调度、符号化变换的程序以及程序特征的公式，使这个过程变得可微。

A2 方法细节

图2展示了Felix编译器的工作流程，其核心方法如下：
- 计算图分区 (§3.1)：Felix首先将输入程序的计算图划分为多个子图，每个子图代表一个或多个张量算子，优化在子图内部进行，子图之间不进行交互优化。
- 符号化调度与程序生成 (§3.2)：对于每个子图，Felix生成多个符号化调度$s_i^*$，每个调度包含一组调度变量$x_i$。这些符号化调度可以生成许多具体的调度。
- 可微性能估计器 (§3.3)：对于每个符号化调度$s_i^*$，Felix创建一个可微的性能估计器$EstmPerf_{i}$。该估计器由两部分组成：
1. 通过分析符号化程序$p_i^*$（由$s_i^*$变换而来），提取出作为调度变量$x_i$的可微函数的一组特征公式$Feat_{p_i^*}(x_i)$。
2. 一个预训练的基于DNN的成本模型 $C$，它将程序特征值映射到预测的执行时间。
- 梯度下降优化 (§3.4)：Felix使用梯度下降法最小化所有符号化调度的目标函数$O(x_i)$，以发现每个子图的高性能调度。
- 整合与生成 (§3.5)：最后，Felix将优化后的子图组合起来，生成一个优化的完整张量程序。

图2. Felix 高层工作流。
图2. Felix 高层工作流。

3.1 计算图分区

  • 分区策略：Felix将输入张量程序的计算图划分为由融合的张量算子组成的子图。这些子图将分轮次进行搜索优化，每轮选择一个子图进行独立优化（详见§3.5）。这种子图分区的思想（源于【38, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI’20】）旨在减小优化的搜索空间，避免对整个程序进行联合优化可能带来的难以处理的复杂性。图分区会以固定的模式融合算子，例如，一个Conv后跟一个ReLU可以融合成一个Conv-ReLU子图，作为一个融合结构进行优化。

3.2 符号化调度与符号化程序生成

  • 符号化调度定义：Felix扩展了Ansor 【38, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI’20】中（具体）的sketch和annotation概念来生成符号化调度。一个sketch是带有未填充可调参数的程序变换列表；Ansor为每个子图生成多个sketch。通过填充这些参数可以产生一个有效的调度。Felix并非使用具体的整数和布尔值来annotate（标注）sketch（这是Ansor唯一能做的方式），而是定义了符号化的调度变量$x_i$并用它们来标注sketch，从而产生符号化的调度$s_i^*$。Felix同时跟踪调度变量值的约束条件$g_{ir}$ ($1 \leq r \leq R_i$)。例如，在图3的调度$s_1^*$中，TILE0满足约束$g_{11}(\text{TILE0}) = (1 \leq \text{TILE0} < K)$。不同符号化程序的约束数量$R_i$各不相同。Felix可以将Ansor生成的任何sketch转换为符号化调度，对于给定的子图，Felix的搜索空间维度与Ansor相同。这些调度变量包括循环分块中的循环大小（图3中的TI$j$, TJ$j$, TK$j$, TILE0）、虚拟线程数（SV0）、循环展开/向量化/并行化的大小，以及计算内联的位置（在TVM 【7, TVM: end-to-end optimization stack for deep learning, 2018, CoRR】中称为ComputeAt变换）。
图3. 在Dense-Add图上生成符号化调度的示例，展示了2个生成的符号化调度及其对应的程序。为了简洁，调度中的一些变换步骤和程序中的长表达式被省略。Felix引入到调度中的变量以粗体显示。
图3. 在Dense-Add图上生成符号化调度的示例，展示了2个生成的符号化调度及其对应的程序。为了简洁，调度中的一些变换步骤和程序中的长表达式被省略。Felix引入到调度中的变量以粗体显示。
  • 符号化程序生成：Felix通过将$s_i^*$中的变换步骤应用于子图的初始程序$p_0$，从每个符号化调度$s_i^*$生成一个符号化程序$p_i^*$，即$p_i^* := T(p_0, s_i^*)$。这些符号化程序包含调度变量$x_i$——循环边界、注解和缓冲区访问索引都是$x_{ij}$的表达式。图3的右列展示了从两个符号化调度$s_1^*$和$s_2^*$派生出的两个程序$p_1^*$和$p_2^*$。

3.3 特征公式提取与重写

  • 特征提取：Felix包含一个程序分析过程，该过程在符号化程序$p_i^*$上运行，以提取多个程序特征，例如程序中加/乘/除操作的总数以及内存缓冲区的重用距离。Felix共使用82个不同的特征来捕捉程序的计算和内存访问特性。由于这些特征依赖于$p_i^*$中包含调度变量$x_i$的循环边界和缓冲区访问索引，因此特征公式是调度变量的函数。下表显示了从Dense-Add图的第一个程序$p_1^*$中提取的几个特征：
    图片描述
    其中，select(c, t, f)是一个分段函数，当布尔表达式c为真时等于t，否则等于f。通常，Felix提取的每个特征表达式只能包含调度变量、常量以及分析过程使用的一组固定操作符，如+、-、*、/、pow、min、max、select等。

  • 提取特征的不可微性：由于某些程序特征的离散性质，提取出的公式包含不连续和不可微的操作符，如select、min和max。一个常见的导致不连续性的场景是，特征的值取决于一个循环层级的边界是否为1（即平凡循环）。int_add特征的公式就表现出这种行为，并包含了select()函数。

  • 表达式可微化：为了获得可微的程序特征公式，Felix使用平滑核为遇到的每个不可微操作符创建平滑的可微近似。Felix中的表达式重写器应用一个内置的重写规则库，每个规则将一个不可微操作符映射到其可微版本。我们通过将不可微函数$f(\cdot)$与一个平滑核$\phi(\cdot)$进行卷积来导出每个近似函数$f_{\text{smooth}}(x)$：
    $$f_{\text{smooth}}(x) := \int_{-\infty}^{x} f(x-t)\phi(t) dt$$
    在这项工作中，我们使用了$\phi(t) := 1 / (1 + t^2)$，相比于其他常见替代方案如高斯核或紧致核，它使得平滑后函数的梯度在数值上更稳定。

  • 可微化示例：经过这个过程，所有生成的函数都将是平滑的（无限可微）。图4比较了两个不可微函数（select(x > 0, 5, 2)和max(x, 0)）与Felix使用的它们的平滑近似。重写规则可以自动地将分析过程产生的所有特征表达式重写为平滑函数的版本。由于不可微操作符的数量很少（少于10个），我们手动导出了平滑版本，这是一次性的工作，且不依赖于具体的变换。这个步骤也可以通过平滑解释（smooth interpretation）等方法【5, Smooth interpretation, 2010, ACM Sigplan Notices】来自动化。
    图4. 两个不可微函数及其各自的平滑可微版本的比较。不可微函数分别是 select(x > 0, 5, 2)（左）和 max(x, 0)（右）。

  • 梯度稳定性：一些程序特征具有指数级增长的表达式和较大的值域。例如，float_add通常是几个变量的乘积，对于普通输入程序，其值可以增长到$10^8 \sim 10^9$。当程序特征取值很大时，一个微小的变化会显得微不足道，导致梯度消失。为了避免梯度消失，Felix对平滑的程序特征取对数，并对每个调度变量$x$执行指数变量替换$x = e^{x'}$，使得$x'$成为新的待优化变量。这两个重写步骤将乘法项转换为加法项，并产生呈线性增长且梯度更稳定的特征表达式。

  • 约束惩罚函数：Felix的表达式重写器还将变量约束$g_{ir}$ ($1 \le r \le R_i$；见§3.2) 重写为可微的约束惩罚函数$p_{ir}(x_i)$ ($1 \le r \le P_i$)。一个约束可以产生一个或多个惩罚函数，即$P_i \ge R_i$。例如，约束$g_{11} = (1 \le \text{TILE0} \le K)$会变成$p_{11}(\text{TILE0}) := 1 - \text{TILE0}$和$p_{12}(\text{TILE0}) := \text{TILE0} - K$。当且仅当所有$p_{ir}(x_i) \le 0$时，调度变量处于有效范围内。对于形如$L \pmod x = 0$的可除性约束（其中$L$是要分块的循环大小，是一个常数），Felix将其替换为一个大小约束$x \le \ln L$，并通过优化后的值舍入来解决可除性问题。Felix不是将$x$舍入到最近的整数，而是将其舍入到最近的$\ln x_j$，其中$x_j$是$L$的一个整数因子。由于$L$是某个张量维度的大小，通常小于$10^6$，因此$x_j$的列表很容易计算。

3.4 使用梯度下降优化调度

  • 目标函数构建：在此阶段，Felix通过组合两个函数为每个符号化程序$p_i^*$创建性能估计器：程序特征$Feat_{p_i^*}(x_i)$和一个预训练的、基于神经网络的成本模型$C$。成本模型$C$将程序特征值作为输入，并输出一个标量作为调度的预测性能，即$EstmPerf_{p_i^*}(x_i) := C(Feat_{p_i^*}(x_i)) \in \mathbb{R}$。然后，Felix使用以下优化问题公式来调整一个子图所有符号化程序的符号变量值：
    $$ \max_{p_i^*} \max_{x_i} C(\text{Feat}_{p_i}(x_i)) \quad \text{s.t.} \quad \forall i, r: q_{ir}(x_i) \leq 0 $$

  • 目标函数的可微化：我们将主要为布尔和整数的变量$x_i$松弛为实值变量，并将上述公式3重写为以下待最小化的目标函数：
    $$O(\mathbf{x}) = \sum_i \left( -C(\text{Feat}_{p_i}(\mathbf{x}_i)) + \lambda \sum_q \max(g_{ir}(\mathbf{x}_i), 0)^2 \right)$$
    这里，我们对$i$求和以同时优化一个子图的所有$x_i$，并将“越高越好”的性能估计器取反，使其变为“越低越好”的目标函数。惩罚项$\max(p_{ir}(x_i), 0)^2$在约束$p_{ir}$未被违反时为0，否则为$p_{ir}(x_i)^2 > 0$，因此最小化此项会推动$x_i$趋向于满足约束。将惩罚项添加到待最小化的函数中，是将有约束的优化问题转化为无约束问题的一种常用方法，例如在DNN权重正则化【11, L2 regularization for learning kernels, 2012, CoRR】和结构化剪枝【35, Learning structured sparsity in deep neural networks, 2016, CoRR】中使用。其中，$\lambda$是控制惩罚函数强度的超参数。目标函数$O(x)$是可微的，因为：(1) $Feat_{p_i^*}(x_i)$通过我们的构造是可微的；(2) $\max(y, 0)^2$函数是可微的，其导数为$2\max(y, 0)$；(3) 每个$p_{ir}(x_i)$通过我们的构造是可微的，因此复合函数$\max(p_{ir}(x_i), 0)^2$也是可微的；最后，(4) 可微函数的和是可微的。

  • 子图调度优化算法：算法1概述了Felix如何使用梯度下降优化一个子图。为了探索更多的搜索空间并避免陷入局部最小值，Felix同时优化多个（nSeeds个）调度。首先，Felix在第10行和第11行提取符号变量、符号调度和目标函数。在第12行，Felix使用拒绝采样法随机采样nSeeds组满足所有约束的调度变量初始值$(\sigma_x)$。第14到19行初始化一个Adam优化器【17, Adam: A method for stochastic optimization, 2014, CoRR】并运行nSteps个优化步骤来最小化公式4给出的目标函数。每一步都评估目标函数$O(x)$在当前调度变量值$\sigma_x$处的梯度，并调用Adam优化器将$\sigma_x$向梯度方向移动。每一步都会更新所有nSeeds个调度。为了获得一个有效的具体调度列表s_valid，Felix接着获取优化过程中遍历的所有变量值，将它们四舍五入为整数，并移除无效的舍入后调度（第20行）。然后，Felix选择最好的nMeasure个调度s_best（按成本模型预测的性能排序），生成它们对应的具体程序，并在目标硬件上进行实证评估。因此，Felix只在目标硬件上进行少量（通常是昂贵的）评估。Felix还使用调度及其测量性能来更新成本模型$C$（这是一个廉价的过程），以便成本模型在下一轮中更好地适应此输入程序的子图。
    算法1：Felix对每个子图的调度调优。

3.5 完整图调优

  • 迭代调优过程：为了优化整个张量程序，Felix首先将程序分解为子图（§3.1），并采用Ansor的基于轮次的调优算法【38, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI’20】来分轮次调优这些子图。在每一轮中，Felix将其子图优化算法（算法1）应用于一个子图。算法2展示了这个迭代调优过程。我们采用Ansor的任务调度器来选择在每一轮中优化哪个子图，但任何其他选择算法也同样适用。最后，在nRounds轮调优之后，Felix为每个子图选择最佳调度，并将它们组合成整个张量程序的完整调度。
    算法2：Felix完整程序调度调优。

3.6 Felix编程接口

  • 易用性：Felix作为一个Python库开发，提供了易于使用的API。图5（原文中）展示了一个使用Felix优化ResNet-50以在Nvidia Xavier NX GPU上运行的示例。关键接口函数在注释中进行了解释。如图所示，使用Felix接口优化一个DNN只需要几行代码。开发者只需提供要优化的网络及其输入形状，并以调优轮次数量的形式指定总调优时间预算。这是现有自动调优工作中的常见做法，因为自动停止标准（例如，收敛时停止）很少被满足。

4. 实现

  • 基础架构：Felix是基于TVM 【7, TVM: end-to-end optimization stack for deep learning, 2018, CoRR】及其Ansor 【38, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI’20】优化器实现的。Felix在TVM TIR中间表示语言中表示其符号化程序，并使用TVM的程序转换过程和硬件后端。它还重用了Ansor的两个算法：图划分算法和sketch生成算法。Felix扩展了Ansor的图划分算法以跟踪创建符号变量的额外信息（§3.1），并处理Ansor的sketch生成算法产生的sketch以生成符号化程序（§3.2）。我们实现了Felix以利用TVM中所有与硬件无关和GPU特定的优化，并支持对Ansor在GPU上调优的所有程序转换参数进行调优。

  • 搜索空间：Felix的搜索空间包括以下可调参数：（1）CUDA线程/块大小，（2）循环分块大小，（3）向量化维度大小，（4）并行化维度大小，以及（5）循环展开因子。这个搜索空间还包括跳过某个转换的选项，因为Felix将大小/因子值为1视为无操作。在Felix和Ansor中，一些优化是无需调优直接应用的；这些决策由TVM和Ansor中的内置规则决定。例如，Felix和Ansor使用基于规则的循环重排，并在适用时贪婪地应用算子融合。对于每个子图，Felix（和Ansor）的搜索空间只包括Ansor的sketch生成算法建议的转换顺序。包含任意转换顺序的搜索超出了本文的范围。

  • 依赖库与实现细节：Felix使用PyTorch 【26, Pytorch: An imperative style, high-performance deep learning library, 2019, Advances in neural information processing systems】进行梯度反向传播和调用算法1中使用的Adam优化器，并使用TenSet 【39, Tenset: A large-scale program performance dataset for learned tensor compilers, 2021, NeurIPS】来定义成本模型和训练成本模型的数据集。我们选择了TenSet提供的多层感知器（MLP）架构，这是一个包含4个线性层和约25万个参数的网络。

  • 代码实现：Felix由13000行C++、Python和Rust代码实现。由于TVM在某些程序转换和程序的某些部分（例如某些循环边界）中需要具体的整数值参数，我们修补了1500行TVM代码以增加对Felix的符号化调度和符号化程序表示的支持。符号化程序生成和特征提取的很大一部分是用C++实现的，而Felix的表达式重写器是用Rust编写的，以便与egg 【36, egg: Fast and extensible equality saturation, 2021, POPL】接口，我们使用这个基于等式饱和的重写框架来实现第3.3节中描述的高效的基于规则的表达式重写。开发者可以通过修改现有TVM转换的源代码以使用符号参数来扩展Felix以调优该转换。此外，重新训练Felix的成本模型和向Felix的表达式重写器添加简化规则是可选的，可以帮助实现更好的搜索结果。

A4 实验环境

  • 数据集/模型：评估使用了6个神经网络：ResNet-50 【15, Deep residual learning for image recognition, 2016, CVPR】, MobileNet-v2 【30, Inverted residuals and linear bottlenecks: Mobile networks for classification, detection and segmentation, 2018, CoRR】, R3D-18 【14, Learning spatiotemporal features with 3d residual networks for action recognition, 2017, CoRR】, DCGAN 【27, Unsupervised representation learning with deep convolutional generative adversarial networks, 2015, CoRR】, Vision Transformer (ViT) 【12, An image is worth 16x16 words: Transformers for image recognition at scale, 2020, CoRR】, 和 LLaMA 【33, Llama: Open and efficient foundation language models, 2023, CoRR】。这些网络涵盖了2D/3D卷积、转置卷积和批量矩阵乘法等常见算子。所有网络都针对float32精度进行推理优化。
  • 模型输入规模：除§6.4外，所有实验的批量大小均为1。对于LLaMA，输入句子长度为100个token。
  • 硬件配置：
    • 服务器GPU: NVIDIA A10G
    • 桌面GPU: NVIDIA RTX A5000
    • 边缘GPU: NVIDIA Xavier-NX
    • 对于Xavier-NX，Felix和TVM的编译过程在另一台配备32核AMD Ryzen 3975WX CPU和512 GB RAM的机器上运行，通过RPC在Xavier-NX上执行张量程序。
  • 软件配置与基线：
    • 代码实现：Felix基于TVM (commit hash: 95aac9224) 和 Ansor 实现，使用C++, Python, Rust。
    • 依赖库：PyTorch 2.2, TensorFlow 2.15, TensorRT 8.6, TenSet。
    • 基线框架：
      1. PyTorch 2.2 (使用TorchInductor后端)。
      2. TensorFlow 2.15 (使用XLA)。
      3. TensorRT 8.6。
      4. Ansor-TenSet：使用在TenSet数据集【39, Tenset: A large-scale program performance dataset for learned tensor compilers, 2021, NeurIPS】上预训练了成本模型的Ansor，这是一个比标准Ansor更强的基线。
  • 搜索参数设置：
    • Felix: 默认使用8个初始调度，200个梯度下降步骤，16次实证测量。每个搜索轮次预测 8 * 200 = 1600 个调度的性能。
    • Ansor: 使用推荐设置，种群大小2048，4代进化，每轮64次实证测量。每个搜索轮次预测 2048 * 4 = 8192 个调度的性能。
  • 成本模型训练：
    • 成本模型在TenSet数据集的子集上训练，该子集包含约25万个调度，覆盖500个子图，包括卷积和线性层等常见瓶颈工作负载。数据集划分为90%训练集和10%验证集。

A4 实验结果

6.1 Felix vs. 现成推理框架

  • 实验内容：比较Felix优化后的6个DNN在3个硬件平台上的推理性能与PyTorch、TensorFlow和TensorRT的性能。
  • 实验结果：
    • 如图6所示，Felix在平均性能上（几何平均值）相比PyTorch提升2.2倍，相比TensorFlow提升1.7倍，相比TensorRT提升1.5倍。在A5000、A10G和Xavier NX上，平均性能提升分别为1.41倍、1.50倍和1.70倍，最高可达4.48倍、5.40倍和10.8倍。
    • 调优时间：如表1所示，Felix平均只需6.9分钟（413秒）的调优时间就能超越性能最好的手动优化库。
  • 分析结论：
    • Felix在包含较小层的网络（如MobileNet-v2和DCGAN）上性能提升尤为显著，因为它能通过更广泛的调度空间搜索来提升并行度和缓存利用率，从而更好地利用GPU计算资源。
    • 性能提升依赖于算子类型。Felix在R3d-18上未超越基线，因为其中的3D卷积在基线框架中被高度优化。但对于其他所有类型的算子，Felix均优于现成框架（见§6.3）。

图6. Felix与深度学习框架（PyTorch, TensorFlow）在三个硬件平台上的DNN推理性能。y轴是某一框架的性能相对于该网络所有框架中最佳性能的归一化值。
表1. Felix在不同DNN和硬件设备上超越性能最佳的手动库所需调优时间（秒）。

6.2 Felix vs. Ansor-TenSet 基于搜索的编译器

  • 实验内容：比较Felix和Ansor-TenSet在不同调优时间内达到的最佳性能。
  • 实验结果：
    • 如图7所示，Felix和Ansor最终收敛到相似的性能水平，但Felix在显著更短的时间内找到更好的调度。
    • 如表2a所示，在批量大小为1时，Felix达到90%峰值性能的搜索时间平均比Ansor-TenSet快5.0倍（A5000）、2.5倍（A10G）和3.2倍（Xavier NX）；达到95%峰值性能的时间平均快3.2倍、1.7倍和4.1倍。
  • 分析结论 (Felix快速收敛的原因)：
    • Felix的梯度下降搜索能更有效地根据其成本模型优化调度，在更少的调优迭代中找到大量预测性能高的调度。
    • 图8显示，Felix的搜索过程能更快地找到大量预测性能高的候选调度（橙色区域），而Ansor的进化搜索（蓝色区域）随机性更强，覆盖范围更广但效率较低。这使得Felix能更快地为硬件评估提供高质量的候选池。

图7. Felix和Ansor-TenSet在调度搜索过程中，最佳性能（显示为网络推理延迟）随搜索时间的变化。
表2. Felix相对于Ansor的调优加速比，通过比较Felix和Ansor达到最佳Ansor性能的90%、95%和99%所需的时间来衡量。(a) 批量大小 = 1
图8. Felix和Ansor中候选调度群体的预测性能随搜索进度的变化。顶行显示调优会话开始时的早期迭代，底行显示接近收敛时的后期迭代（两个工具的线重叠）。

6.3 单个算子性能比较

  • 实验内容：在RTX A5000上比较Felix、Ansor、PyTorch和TensorFlow对不同类型张量算子的性能。
  • 实验结果：如图9所示，Felix在8种算子类型中的7种上优于PyTorch和TensorFlow，并且性能与Ansor相当。唯一的例外是3D卷积，PyTorch和TensorFlow的性能优于Felix和Ansor。
  • 分析结论：对于常见且计算密集的算子（如3D卷积），手动调优库投入了大量精力进行优化，因此性能极高。然而，这些库的性能在不同硬件间差异很大，需要针对性调优。相比之下，Felix能够自动进行硬件特定的调优，无需手动干预。
图9. Felix、Ansor和手动优化库（PyTorch、TensorFlow）在RTX A5000上的单个算子性能。y轴是某一框架的性能相对于该算子所有框架中最佳性能的归一化值。
图9. Felix、Ansor和手动优化库（PyTorch、TensorFlow）在RTX A5000上的单个算子性能。y轴是某一框架的性能相对于该算子所有框架中最佳性能的归一化值。

6.4 不同批量大小下的性能

  • 实验内容：在批量大小为16的场景下，比较Felix和Ansor-TenSet在RTX A5000上的调优效率。
  • 实验结果：如图10和表2b所示，在此设置下，Felix达到90%峰值性能的搜索时间平均比Ansor-TenSet快5.8倍，达到95%峰值性能的时间快4.9倍。
  • 分析结论：即使在输入批量大小增加的情况下，Felix依然能快速收敛到高性能的调度。

表2b. 批量大小 = 16
图10. Felix和Ansor-TenSet的调度搜索过程中最佳性能（网络推理延迟）随搜索时间的变化。网络在RTX A5000上针对输入批量大小为16进行优化。

A7 相关工作

  • 自动张量程序生成：许多张量程序编译器，如Halide【28, Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines, 2013, PLDI】、TVM【7, TVM: end-to-end optimization stack for deep learning, 2018, CoRR】和Tiramisu【4, Tiramisu: A polyhedral compiler for expressing fast and portable code, 2019, CGO】，都提供了将计算与调度分离的能力。这些编译器通过自动调优来搜索好的调度。Halide有多个基于不同调优技术的自动调度器【2, Learning to optimize halide with tree search and random programs, 2019, TOG】、【21, Differentiable programming for image processing and deep learning in halide, 2018, TOG】、【25, Automatically scheduling halide image processing pipelines, 2016, TOG】。TVM也有AutoTVM【7, TVM: end-to-end optimization stack for deep learning, 2018, CoRR】、Ansor【38, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI’20】和MetaSchedule【31, Tensor program optimization with probabilistic programs, 2022, NeurIPS】等调优系统。这些系统都使用组合离散搜索算法，如遗传算法、树搜索或在有限空间内进行暴力搜索。与这些工作不同，Felix直接在张量程序调度上应用梯度下降作为搜索技术，利用了梯度下降对平滑函数的高效性。本文将Ansor【38, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI’20】作为基准进行比较，因为它在TVM生态系统中性能优于并取代了AutoTVM。Felix的方法与现有编译器是互补的，其基于梯度的搜索技术也可应用于其他优化系统。

  • 基于梯度的DNN搜索问题优化：MindMappings【16, Mind mappings: enabling efficient algorithm-accelerator mapping space search, 2021, ASPLOS】引入了一种基于梯度的算法，用于搜索定制加速器的映射空间，通过梯度下降优化数据分块大小等调度参数。但MindMappings为每种算子类型都需要收集新数据集并训练独立的成本模型，且其搜索空间针对可编程加速器，不适用于通用GPU/CPU。DARTS【22, DARTS: differentiable architecture search, 2018, CoRR】和SSL【35, Learning structured sparsity in deep neural networks, 2016, CoRR】等工作通过将离散优化问题松弛为连续问题并应用梯度技术来解决神经架构搜索和结构化剪枝问题。然而，这些技术的目标函数通常已经是分析表达式，只需手动进行数学重构即可微分。Felix的独特性在于它能自动为程序特征创建符号表达式，并将其转换为可微的对应物以进行梯度下降优化。

  • 自动微分：自动微分（AD）是一种为给定程序自动生成梯度的技术，广泛应用于PyTorch【26, Pytorch: An imperative style, high-performance deep learning library, 2019, Advances in neural information processing systems】、TensorFlow【1, Tensorflow: Large-scale machine learning on heterogeneous distributed systems, 2016, CoRR】等深度学习框架。Enzyme【24, Instead of rewriting foreign code for machine learning, automatically synthesize fast gradients, 2020, NeurIPS】等通用AD框架甚至可以在LLVM IR级别推导梯度。这些AD框架的目标是为给定程序（其本身代表一个目标函数）生成其输入的梯度。这与Felix的目标完全不同，Felix是为了优化给定程序的执行性能而生成其调度参数的梯度。

A5 结论

本文介绍了Felix，一个新颖的、基于梯度的张量程序编译器优化框架。我们展示了如何创建程序转换调度的可微空间，并利用梯度下降来寻找高性能的调度。Felix优化算法的新颖之处在于对调度空间进行了审慎的连续松弛，并生成了一个适合梯度下降优化的可微成本函数。

评估结果表明，Felix中高效的基于梯度的搜索能够比PyTorch、TensorFlow和TensorRT提供的程序更快地找到更优的张量程序。此外，我们展示了Felix相对于最先进的进化搜索的优势：它生成高性能程序的速度明显快于TVM Ansor。我们的评估证明，Felix可以广泛应用于从服务器、桌面到资源受限的边缘设备的各种商用GPU。

我们相信Felix背后的核心思想——通过变量松弛和数值优化来优化计算图——是通用的，并且可以基于其他操作计算图的张量编译器（如Halide【28, Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines, 2013, PLDI】、MLIR【20, MLIR: A compiler infrastructure for the end of moore’s law, 2020, CoRR】、TACO【18, The tensor algebra compiler, 2017, OOPSLA】、Tiramisu【4, Tiramisu: A polyhedral compiler for expressing fast and portable code, 2019, CGO】）来实现。我们视此方法为将基于梯度的优化推广到通用编程语言中具有更复杂成本特征和控制流的更广泛程序类别的第一步。

方法细节中的引用汇总

  • 【1, Tensorflow: Large-scale machine learning on heterogeneous distributed systems, 2016, CoRR】: 在第7节（相关工作）中引用，作为提供自动微分能力的深度学习框架示例。
  • 【2, Learning to optimize halide with tree search and random programs, 2019, TOG】: 在第2节（问题陈述）和第7节（相关工作）中引用，作为基于搜索的张量编译器（Halide自动调度器）的示例，该调度器使用树搜索。
  • 【3, A deep learning based cost model for automatic code optimization, 2021, CoRR】: 在第2节（问题陈述）中引用，作为基于搜索的张量编译器（Tiramisu）的示例。
  • 【4, Tiramisu: A polyhedral compiler for expressing fast and portable code, 2019, CGO】: 在第1节（引言）、第7节（相关工作）和第8节（结论）中引用，作为提供计算与调度分离能力的张量程序编译器示例，其思想可以与Felix结合。
  • 【5, Smooth interpretation, 2010, ACM Sigplan Notices】: 在第3.3节（特征公式提取与重写）中引用，作为一种可以自动化生成平滑可微函数的方法。
  • 【6, Mxnet: A flexible and efficient machine learning library for heterogeneous distributed systems, 2015, CoRR】: 在第7节（相关工作）中引用，作为提供自动微分能力的深度学习框架示例。
  • 【7, TVM: end-to-end optimization stack for deep learning, 2018, CoRR】: 在第1节（引言）、第2节（问题陈述）、第3.2节（符号化调度与程序生成）、第4节（实现）和第7节（相关工作）中多次引用，作为Felix实现的基础框架，并提及其ComputeAt变换和AutoTVM调优系统。
  • 【11, L2 regularization for learning kernels, 2012, CoRR】: 在第3.4节（使用梯度下降优化调度）中引用，作为将惩罚项加入目标函数以处理约束的类似方法（DNN权重正则化）。
  • 【16, Mind mappings: enabling efficient algorithm-accelerator mapping space search, 2021, ASPLOS】: 在第7节（相关工作）中引用，作为另一个使用梯度下降进行调度优化的工作，并与Felix进行对比，指出MindMappings的局限性。
  • 【17, Adam: A method for stochastic optimization, 2014, CoRR】: 在第1节（引言）和第3.4节（使用梯度下降优化调度）中引用，作为Felix中使用的具体梯度下降优化器。
  • 【18, The tensor algebra compiler, 2017, OOPSLA】: 在第8节（结论）中引用，作为Felix思想可以应用的另一种张量编译器。
  • 【20, MLIR: A compiler infrastructure for the end of moore’s law, 2020, CoRR】: 在第8节（结论）中引用，作为Felix思想可以应用的另一种张量编译器。
  • 【21, Differentiable programming for image processing and deep learning in halide, 2018, TOG】: 在第7节（相关工作）中引用，作为Halide自动调度器的一个示例，并提及其为图像处理流水线自动推导梯度的能力。
  • 【22, DARTS: differentiable architecture search, 2018, CoRR】: 在第1节（引言）和第7节（相关工作）中引用，作为将离散搜索问题松弛为连续可微问题并应用梯度下降的代表性工作（神经架构搜索）。
  • 【24, Instead of rewriting foreign code for machine learning, automatically synthesize fast gradients, 2020, NeurIPS】: 在第7节（相关工作）中引用，作为通用的自动微分框架（Enzyme），并与Felix的目标进行区分。
  • 【25, Automatically scheduling halide image processing pipelines, 2016, TOG】: 在第7节（相关工作）中引用，作为Halide自动调度器的一个示例，该调度器在有限空间内进行暴力搜索。
  • 【26, Pytorch: An imperative style, high-performance deep learning library, 2019, Advances in neural information processing systems】: 在第1节（引言）、第4节（实现）和第7节（相关工作）中引用，作为现有深度学习框架和Felix实现中使用的梯度反向传播库。
  • 【28, Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines, 2013, PLDI】: 在第1节（引言）、第7节（相关工作）和第8节（结论）中引用，作为重要的张量程序编译器，其思想可与Felix结合。
  • 【31, Tensor program optimization with probabilistic programs, 2022, NeurIPS】: 在第7节（相关工作）中引用，作为TVM的MetaSchedule调优系统，其使用遗传算法，与Felix的方法正交。
  • 【35, Learning structured sparsity in deep neural networks, 2016, CoRR】: 在第1节（引言）、第3.4节（使用梯度下降优化调度）和第7节（相关工作）中引用，作为将离散搜索问题松弛为连续可微问题并应用梯度下降的代表性工作（结构化剪枝）。
  • 【36, egg: Fast and extensible equality saturation, 2021, POPL】: 在第4节（实现）中引用，作为Felix中用于高效表达式重写的框架。
  • 【38, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI’20】: 在第1节（引言）、第3.1节（计算图分区）、第3.2节（符号化调度与程序生成）、第3.5节（完整图调优）、第4节（实现）和第7节（相关工作）中多次引用，作为Felix实现的基础、比较的基线，并借鉴了其图分区和sketch生成的思想。
  • 【39, Tenset: A large-scale program performance dataset for learned tensor compilers, 2021, NeurIPS】: 在第1节（引言）、第4节（实现）、第5节（实验方法）和第6.2节（实验结果）中引用，作为Ansor-TenSet基线和Felix成本模型所使用的性能数据集。
  • 【40, Neuron-level structured pruning using polarization regularizer, 2020, NeurIPS】: 在第3.4节（使用梯度下降优化调度）中引用，作为结构化剪枝的另一个例子。