Optimal Kernel Orchestration for Tensor Programs with Korch

论文标题：利用Korch实现张量程序的最优核函数编排

作者/机构:
Muyan Hu (伊利诺伊大学厄巴纳-香槟分校), Ashwin Venkatram (AMD), Shreyashri Biswas (卡内基梅隆大学), Balamurugan Marimuthu (Sambanova Systems), Bohan Hou (卡内基梅隆大学), Gabriele Oliaro (卡内基梅隆大学), Haojie Wang (清华大学), Liyan Zheng (清华大学), Xupeng Miao (卡内基梅隆大学), Jidong Zhai (清华大学), Zhihao Jia (卡内基梅隆大学)

A1 主要贡献

本文旨在解决深度神经网络（DNN）在现代硬件平台上执行时的核函数编排（Kernel Orchestration）问题。核函数编排指的是将DNN中不同算子定义的计算映射到GPU核函数执行的任务。

核心问题：
现有的深度学习框架通过贪婪地应用算子融合（operator fusion）来优化核函数编排。这种方法将多个算子的计算融合成一个单一的核函数，以减少核函数启动开销和不必要的中间结果实例化。然而，这种方法存在两大局限性：
1. 算子级别的融合粒度过粗： 某些算子（如Softmax）在单一核函数中执行时性能不佳，因为其内部包含不同并行度和内存访问模式的子计算。将这些算子分解为更细粒度的组件，并跨算子融合具有相同特性的组件，可以带来性能提升，但现有框架无法表示这类优化。
$$ \text{softmax}(x_i) = e^{x_i} / \sum_j e^{x_j} $$
2. 依赖手动设计的贪婪规则： 现有的融合方法依赖于专家手动设计的规则，这需要大量的工程努力，并且可能错过一些难以手动发现的微妙优化。此外，为一种设备设计的规则不一定适用于其他类型的设备。

研究目标与创新点：
为解决上述问题，本文提出了Korch，一个张量程序优化器，旨在发现张量程序的最优核函数编排策略。其主要贡献如下：
* 识别核函数编排任务： 明确提出了“核函数编排”这一任务，并揭示了现有方法中缺失的一系列核函数优化机会。
* 引入算子分裂（Operator Fission）： 与直接融合算子不同，Korch首先将张量算子分解（fission）为一组基本的张量代数原语（primitives）。这种分解使得跨算子的细粒度优化成为可能，从而最大化核函数效率。
* 基于二元线性规划的最优映射： Korch将核函数编排问题形式化为一个约束优化问题，并利用现成的二元线性规划（Binary Linear Programming, BLP）求解器来发现最优的编排策略，即如何将原语图映射到核函数。
* 实现并评估Korch系统： 实现了Korch系统，并在多种DNN上进行了全面评估，结果显示在NVIDIA V100 GPU上性能提升高达1.7倍，在A100 GPU上高达1.6倍。

A2 Korch系统概述


图 1. Korch概览。

Korch工作流程概述
Korch的输入是一个待优化的张量程序，表示为一个计算图，其中节点是张量代数算子，边是它们之间的数据依赖关系。与先前的工作【【19，Optimizing dnn computation with relaxed graph substitutions，SysML’19】，【35，Pet: Optimizing tensor programs with partially equivalent transformations and automated corrections，OSDI 21】】类似，Korch首先将输入的计算图划分为较小的子图，以减少每个子图的优化空间，同时保留优化机会。

算子分裂与图优化
对于每个子图，Korch的算子分裂引擎（operator fission engine）使用一种基于规则的方法，将每个算子解耦为一个或多个张量代数原语，并生成一个功能上与原始计算图等价的原语图。Korch的原语图优化器采用了先前工作【【18，Taso: optimizing deep learning computation with automatic generation of graph substitutions，SOSP'19】】中引入的超优化技术，并使用由TASO【【18，Taso: optimizing deep learning computation with automatic generation of graph substitutions，SOSP'19】】发现的图变换来优化原语图。通过将张量算子解耦为细粒度的原语，Korch能够实现额外的原语图级别优化，这些优化在算子图级别是不可行的（详见§3）。这种分解还允许Korch通过直接选择原语图的子图作为候选核函数，从而执行更灵活的核函数编排（详见§4）。

核识别与编排
为了将张量代数原语映射到GPU核函数，Korch的核函数识别器（kernel identifier）使用深度优先搜索算法来识别给定原语图的所有可能核函数。然后，Korch的核函数分析器（kernel profiler）测量每个候选核函数的运行时性能。Korch将核函数编排形式化为一个二元线性规划（BLP）问题，并使用一个现成的BLP求解器来发现最优的核函数编排策略。具体来说，Korch的核函数编排优化器将所有已识别的核函数及其执行延迟作为输入，并采用BLP求解器来选择一个最优策略，该策略执行候选核函数的一个子集，并根据输入计算张量程序的输出，同时最小化总执行成本。

可执行文件生成
最后，最优执行策略被Korch的可执行文件生成器（executable generator）用来为输入的张量程序生成一个可执行文件。我们将在§3和§4中分别描述算子分裂和核函数编排组件。

A3 方法细节

3. 算子分裂（Operator Fission）

算子级融合的局限性
现有的深度学习框架通过应用算子融合（operator fusion）来优化核函数编排，机会性地将多个张量算子融合成一个单一的GPU核函数【【7，TVM: end-to-end optimization stack for deep learning，2018】，【8，Learning to optimize tensor programs，NeurIPS'18】】。然而，直接在算子级别应用融合会错失各种细粒度的、跨算子的优化机会。这是因为张量算子是根据其代数语义和属性设计的，而不是根据其在GPU上的实现。因此，对于具有复杂数学语义的算子，如归一化和聚合，在单个核函数内执行该算子定义的所有计算会导致次优性能。

基于算子分裂的原语分解
基于上述观察，Korch首先应用算子分裂（operator fission）将每个张量算子解耦为一小组张量代数原语（tensor algebra primitives），每个原语都执行涉及统一并行度和数据访问模式的基本代数计算。Korch考虑的原语根据每个原语的输入和输出元素之间的关系分为四类。对于每个原语$p$，我们用$O(I)$表示在$n$个输入张量$I = (I_1, ..., I_n)$上运行$p$的输出张量。对于一个$m$维张量$T$，我们用$T[\vec{x}]$表示在位置$\vec{x} = (x_1, ..., x_m)$的输出值。类似地，$I_k[\vec{x}]$表示第$k$个输入张量$I_k$在位置$\vec{x}$的输入值。我们使用上述符号来介绍Korch考虑的四类张量代数原语。

Elementwise（元素级）原语
对于一个元素级原语$p$，输出张量与所有输入张量具有相同的形状和数据布局，并且每个输出元素的计算依赖于相同位置的输入元素：
$O[\vec{x}] = f(I_1[\vec{x}], I_2[\vec{x}], ..., I_n[\vec{x}])$
元素级原语不涉及布局转换，通常可以作为预处理或后处理步骤与其他原语融合。

Reduce（规约）和 Broadcast（广播）原语
一个规约原语接收一个张量作为输入，并计算输入张量每一行沿给定维度的聚合结果：
$$O[x_1, \dots, x_{k-1}, x_{k+1}, \dots, x_m] = \bigoplus_{x_k} I_1[x_1, \dots, x_{k-1}, x_k, x_{k+1}, \dots, x_m],$$
该公式沿$I_1$的第$k$维进行聚合，其中⊕是聚合器。一个广播原语执行相反的操作，接收一个输入张量并沿给定维度复制该输入张量：
$$O[x_1, \dots, x_{k-1}, x_k, x_{k+1}, \dots, x_m] \\ = I_1[x_1, \dots, x_{k-1}, x_{k+1}, \dots, x_m].$$

Layout Transformation（布局转换）原语
一个布局转换原语接收一个张量作为输入，并转换其数据布局而不执行算术运算：
$O[\vec{x}] = I_1[L(\vec{x})],$
其中$\vec{x'} = L(\vec{x})$是从输出张量的位置$\vec{x}$到输入张量的位置$\vec{x'}$的一对一映射。注意，一个布局转换原语可以有多个输入/输出张量。例如，拼接（concatenation）和分割（split）被认为是布局转换原语，因为它们需要改变输入张量的底层布局。

Linear Transformation（线性转换）原语
最后，Korch使用线性转换原语来捕获DNN中的计算密集型算子，如矩阵乘法和卷积。具体来说，如果一个原语$p$的输出对所有输入张量都是线性的，那么它就被认为是一个线性转换原语：
$$\forall Y, Z, \vec{x}, 1 \leq k \leq n . \quad O(I_1, \ldots, I_{k-1}, Y+Z, \ldots)[\vec{x}]$$
$$= O(I_1, \ldots, I_{k-1}, Y, \ldots)[\vec{x}] + O(I_1, \ldots, I_{k-1}, Z, \ldots)[\vec{x}],$$
$$\alpha \cdot O(I_1, \ldots, I_{k-1}, Y, \ldots)[\vec{x}] = O(I_1, \ldots, I_{k-1}, \alpha \cdot Y, \ldots)[\vec{x}]$$

算子分裂规则示例
表1显示了四种原语类型的代表性算子。对于每个DNN算子，Korch要求开发者指定一个算子分裂规则，该规则用于将DNN算子解耦为指定的原语。每个算子分裂规则都是从一个DNN算子到功能等价的原语图的转换。为了展示设计算子分裂规则的复杂性，图3显示了Softmax（见公式(1)）的算子分裂规则，它被分解为一个元素级指数运算、一个规约、一个广播和一个元素级除法原语。
表 1. 一些常用原语的示例。

图 3. Softmax的算子分裂规则。

算子分裂的第一个优势：图优化
算子分裂使得两类在现有深度学习框架中缺失的重要优化成为可能。首先，通过将算子解耦为细粒度的原语，算子分裂使得在原语图上进行优化转换成为可能；这些转换在算子级别是无法应用的。图2b展示了这样一个优化Transformer【【34，Attention is all you need，NeurIPS'17】】中多头注意力的例子。首先，通过引入一个元素全为1的常数输入张量$I_{ones}$，Softmax的规约原语被替换为一个线性转换原语（即MatMul）。其次，通过应用先前工作【【18，Taso: optimizing deep learning computation with automatic generation of graph substitutions，SOSP'19】】发现的变换，Softmax的元素级除法原语与随后的MatMul交换了位置。最后，通过引入一个Pad和一个Split原语，将两个具有共享输入的MatMul融合在一起。这些变换使得Korch能够将Softmax中的规约原语与一个MatMul融合。

算子分裂的第二个优势：灵活的核编排
算子分裂的第二个优势是实现了更灵活的核编排。具体来说，对于图2c中的优化后原语图，Korch将Div和Exp映射到核函数①，其性能与在单个核函数中仅执行Div相同，这是由于元素级计算的算术强度较低。类似地，Transpose和Pad被映射到核函数②，其性能与单个Transpose核函数相同。MatMul本身被映射到核函数③，因为它计算密集，将其映射为单个核函数允许Korch使用CUDA库提供的核函数（详见§5.2）。最后，Split、Broadcast和Div原语都映射到核函数④。Split和Broadcast都不引入算术运算，执行成本低廉。此优化的性能提升是由于核函数④比启动一个Softmax核函数快得多。
图 2. 算子分裂使得后续在原语图上的优化转换成为可能。在图2b中，相同颜色的虚线框表示对原语图的一次转换。这三次转换的组合将Softmax中的规约原语和随后的MatMul融合成一个单一的MatMul。

新算子支持
上面列出的四类原语足以表示当今DNN中使用的各种张量算子，包括我们评估中的所有DNN基准测试。然而，有些张量算子无法使用上述原语表示。一个这样的算子是Topk，它返回输入张量沿给定维度的$k$个最大元素。对于这类算子，Korch将它们视为不透明原语，并且只会在图优化中优化原语图的其余部分。

4. 核函数编排（Kernel Orchestration）

核编排的目标与约束
本节描述Korch的核函数编排引擎，它以一个原语图作为输入，并发现一个最优策略来将原语映射到核函数，同时最小化执行延迟。一个原语图表示为一个有向无环图$G = (P, E)$，其中每个节点$p \in P$是一个张量代数原语，每条边$(p_1, p_2) \in E$表示一个张量，该张量是$p_1$的输出和$p_2$的输入。GPU上的计算被组织成核函数，每个核函数是一个以单程序多数据（SPMD）方式由多个线程同时执行的函数。每个核函数执行一个或多个原语。为了避免核函数之间的循环依赖（即两个核函数相互依赖才能开始），一个核函数中包含的原语必须形成$G$的一个凸子图（convex subgraph）。

凸子图定义
定义1 (凸子图): 对于一个原语图 $G = (P, E)$，一个节点集 $P' \subseteq P$ 构成 $G$ 的一个凸子图，当且仅当不存在节点 $p_1, p_2 \in P'$ 和另一个节点 $p \in P \setminus P'$ 使得 $p_1 \rightsquigarrow_G p$ 且 $p \rightsquigarrow_G p_2$，其中 $p \rightsquigarrow_G p'$ 表示在 $G$ 中存在从 $p$ 到 $p'$ 的路径。

执行状态与凸子图识别
对于图4a中的原语图，节点集 $\{p_1, p_4, p_8\}$ 构成$G$的一个凸子图，因此可以在一个核函数中执行。另一方面，我们不能在单个核函数中执行节点集 $\{p_1, p_2, p_5\}$（$G$的一个非凸子图），因为它与执行$p_3$的核函数存在循环依赖（即$p_5$依赖于$p_3$，而$p_3$又依赖于$p_1$）。Korch使用执行状态来识别给定原语图的所有凸子图。

执行状态定义
定义2 (执行状态): 一个节点集 $P' \subseteq P$ 是原语图 $G = (P, E)$ 的一个执行状态，如果对于任何边 $(p_1, p_2) \in E$，$p_2 \in P'$ 蕴含 $p_1 \in P'$。

定理1：凸子图与执行状态的关系
从一个执行状态转换到另一个执行状态使得Korch能够识别所有可能的凸子图，如定理1所示。$G$的每个凸子图的计算可以在一个候选核函数中执行。
定理1: 一个节点集 $P' \subseteq P$ 构成 $G = (P, E)$ 的一个凸子图，当且仅当 $P'$ 是两个执行状态 $P_1$ 和 $P_2$ 的差集：$P' = P_1 \setminus P_2$。

定理1证明
证明: 如果 $P' = P_1 \setminus P_2$，下面将证明$P'$是一个凸子图。假设存在$p_1, p_2 \in P'$和$p \in P \setminus P'$使得$p_1 \rightsquigarrow_G p$和$p \rightsquigarrow_G p_2$。因此，$p_1, p_2 \in P_1$且$p_1, p_2 \notin P_2$。同时我们可得$p \notin P_1$或者$p \in P_2$。如果$p \notin P_1$，那么$p_1 \rightsquigarrow p$和$p_1 \in P_1$将与$P_1$是执行状态相矛盾。如果$p \in P_2$，那么$p_1 \rightsquigarrow p$和$p_1 \notin P_2$将与$P_2$是执行状态相矛盾。综上，假设是错误的，所以$P'$是一个凸子图。
如果$P'$是凸图，令$S = \{p | p \notin P' \text{ and } \exists(p, p') \in E, p' \in P'\}$。那么我们可以构造两个执行状态$P_1 = \{p | p \in P' \text{ or } \exists p' \in P_1, (p, p') \in E\}$和$P_2 = \{p | p \in S \text{ or } \exists p' \in P_2, (p, p') \in E\}$。从$P'$的凸性我们可以得到$P' = P_1 \setminus P_2$。□

执行状态的数量
直观上，执行状态的数量随原语图的深度（即关键路径上的算子数量）线性增加，随宽度（即并发算子数量）指数增加。虽然一个原语图的执行状态数量可能随原语数量指数增长，但我们观察到现代DNN架构通常是深而不是宽的，导致执行状态数量有限，从而使得显式考虑所有执行状态是可行的。为简单起见，我们预计一个原语图大约有 $O(|P|)$ 个执行状态和 $O(|P|^2)$ 个候选核函数，其中 $|P|$ 是图中原语的数量。这种简化在先前的工作中也被采用【【31，Piper: Multidimensional planner for dnn parallelization，NeurIPS'21】】。

核识别器概述
接下来，我们描述Korch的核识别器，它使用深度优先搜索算法来识别原语图中的所有有效核函数。

4.1 核函数识别器（Kernel Identifier）

算法1概述
算法1展示了在一个原语图$G$中识别所有可能核函数的算法。该算法从一个空的执行状态$X = \emptyset$开始，并使用深度优先搜索（DFS）来枚举所有可能的执行状态，这些状态被维护在一个数据库$B$中。对于每对执行状态$(D_1, D_2)$，它们的集合差识别出一组原语$P' = D_2 \setminus D_1$，这些原语可能构成一个核函数。
算法 1. 深度优先搜索以识别核函数。Profiling函数接收一组原语作为输入，为定义的张量代数计算生成一个核函数，并分析该核函数的运行时性能。如果无法生成这样的核函数，Profiling返回∞。我们将在§5.2中描述核函数分析器。

核的输入与输出
考虑一个子图$G' = (P', E')$，我们现在识别$G'$的输入和输出。显然，$G'$的输入应该是$G'$中所有入度为零的节点。然而，$G'$的输出并不仅仅是所有出度为零的节点，因为一些出度为正的节点可能会成为另一个候选核函数子图的输入。我们如下定义$G'$的可能输出集：

可能输出集定义
定义3 (可能输出集): 一个节点集$O$是候选子图$G' = (P', E')$的一个可能输出集，当且仅当 $\forall o \in O, \exists (o, p) \in E$ 使得 $p \in P \setminus P'$。

核性能分析
$P'$和$O$都被发送到核函数分析器（kernel profiler）为其原语生成一个核函数。如果能成功生成这样的核函数，核函数分析器还会返回测量的该核函数的执行延迟；否则，分析器返回无穷大，以表明$P'$不被Korch的DNN后端支持。§5.2介绍了Korch用于为给定原语集生成和分析核函数的核函数分析器。对于图4a中的原语图，图4b显示了由Korch的核函数分析器识别出的所有有效核函数。
图 4. 自注意力【【34，Attention is all you need，NeurIPS'17】】子图的核函数编排示例。图4a中的阴影子图与图2中的原语图转换结果相似。

4.2 核函数编排优化器（Kernel Orchestration Optimizer）

核编排问题定义
本节首先定义核函数编排问题，然后将其形式化为一个约束优化任务。Korch使用二元整数线性规划算法来发现一个最小化GPU上总执行成本的核函数编排策略。对于给定的原语图$G$，令$k_1, ..., k_M$表示由核函数识别器（见§4.1）发现的所有核函数。核函数编排优化器接收一个原语图$G$和所有已识别的核函数作为输入，并发现一个最优的核函数执行策略，该策略由一个$M$维二元向量$\vec{u} = \langle u_1, ..., u_M \rangle$表示，其中如果$k_i$在该策略中被执行，则$u_i = 1$，否则$u_i = 0$。Korch遵循先前工作【【18，Taso: optimizing deep learning computation with automatic generation of graph substitutions，SOSP'19】】的做法，假设一个执行策略的运行时间是各个核函数运行时间的总和：
$$Cost(\vec{u}) = \sum_{i=1}^{M} c_i u_i,$$
其中$c_i$是核函数$k_i$的测量运行时间。

BLP方法与冗余计算
为了发现一个最小化运行时间的核函数执行策略，我们的关键思想是将核函数编排形式化为一个二元线性规划（BLP）任务，并使用一个现成的BLP求解器找到一个最优策略，这将在§5中介绍。Korch使用两个矩阵$I$和$O$来表示核函数和原语之间的相关性。具体来说，$I$和$O$都是大小为$M \times |P|$的二元矩阵，其中$I_{ij} = 1$当且仅当原语$p_j$是核函数$k_i$的输入，而$O_{ij} = 1$当且仅当原语$p_j$是核函数$k_i$的输出。因此，$\sum_{i=1}^{M} O_{ij} u_i$计算了原语$p_j$被执行的次数。Korch与先前工作的一个关键区别在于，现有的张量程序优化器以不相交的方式（即没有冗余计算）将DNN架构中定义的计算划分为核函数，而Korch允许一个原语被多次执行（即$\sum_{i=1}^{M} O_{ij} u_i > 1$）。这一放宽是基于一个重要观察：现代GPU的浮点吞吐量增长速度远高于其内存带宽，如图5所示。因此，Korch允许原语被多次执行，以便机会性地减少内存访问和核函数启动开销。例如，图4c展示了一个最优的核函数编排策略，该策略执行$p_1$三次以减少核函数启动和数据移动开销。
图 5. 比较几代GPU的内存带宽和浮点吞吐量。y轴是相对于P100 GPU性能的归一化值。

与张量重物质化的比较
张量重物质化（Tensor rematerialization）是DNN训练中常用的一种技术，通过在前向传播过程中丢弃一些中间结果，并在反向传播过程中重新物质化（re-materializing）这些被丢弃的张量，来减少DNN训练的内存开销。张量重物质化以在一次训练迭代中多次重新执行某些核函数为代价，降低了DNN训练的峰值内存使用量。与减少对GPU设备内存的总体访问不同，张量重物质化实际上增加了内存访问，因为某些核函数被多次执行。相比之下，Korch中的原语冗余旨在允许对具有重叠原语的核函数进行机会性融合。

BLP约束：输出约束
我们现在描述我们的二元线性规划公式。对于给定的原语图$G = (P, E)$，令$T \subseteq P$表示输出原语的集合。BLP任务通过考虑以下约束来最小化公式(2)中定义的目标函数。输出约束保证了一个有效策略会执行所有作为DNN模型输出的原语：
$$\sum_{i=1}^{M} O_{i j} u_{i} \geq 1 \quad \forall p_{j} \in \mathcal{T}.$$

BLP约束：依赖约束
依赖约束检查不同核函数之间的数据依赖关系——一个核函数只有在其所有输入张量都已被先前的核函数计算出来后才能执行：
$$\sum_{i=1}^{M} O_{ij} u_i \geq I_{k,j} u_k \quad \forall p_j \in \mathcal{P}, k \in [1, M]$$
注意，Korch不要求所有原语都被执行，我们的公式隐式地执行了剪枝优化。Korch的核函数执行策略可以看作是用候选核函数子图覆盖整个计算图，仅受子图依赖关系的约束。

5. 实现

5.1 算子分裂
基于ONNX的实现
我们的实现使用ONNX格式【【24，ONNX: Open neural network exchange，2022】】来表示原语图，该格式被广泛使用且在数学上是完备的。特别是，ONNX算子【【25，ONNX Operators，2022】】包括了§3中介绍的所有原语类型。Korch的算子分裂引擎接收一个ONNX图作为输入，并执行基于规则的算子分裂来构建一个等价的原语图。算子分裂引擎的输出也是ONNX格式。

5.2 核函数编排优化器
核函数分析器
Korch的核函数分析器接收一个原语图作为输入，为输入图生成一个GPU核函数，并分析该核函数的运行时性能。Korch遵循先前工作的做法，要求每个原语图只有一个输出【【7，TVM: end-to-end optimization stack for deep learning，2018】】。因此，每个候选核函数都有任意数量的输入张量并产生一个输出张量。对于每个候选核函数，分析器检查子图是否包含任何线性转换原语。如果没有找到线性转换原语，分析器将当前候选核函数识别为内存密集型，并将子图发送到TVM的MetaSchedule进行超参数调优。否则，分析器将当前候选核函数识别为计算密集型。

MetaScheduler
MetaScheduler【【29，Tensor program optimization with probabilistic programs，NeurIPS'22】】是在TensorIR中开发的一个概率性调度器DSL，它统一了先前存在的方法（即AutoTVM【【8，Learning to optimize tensor programs，NeurIPS'18】】和Ansor【【37，Ansor: generating high-performance tensor programs for deep learning，OSDI'20】】）。Korch使用MetaScheduler为给定的原语图和目标硬件生成高性能核函数。具体来说，Korch根据预定义的调度规则将ONNX图降级为TensorIR调度，并利用Ansor后端生成性能最佳的调度。然后，核函数分析器测量生成核函数的性能，并将测量的延迟、降级的调度器和生成的CUDA代码返回给核函数编排优化器。由于内存密集型核函数的调度相对简单，大多数可以通过MetaScheduler在2分钟内完成调优。

计算密集型算子
对于计算密集型算子，Korch遵循先前工作（例如EinNet【【38，{EINNET}: Optimizing tensor programs with {Derivation-Based} transformations，OSDI'23】】）使用的启发式方法，直接将这些原语图降级到供应商库（例如cuDNN和cuBLAS），而不是TVM。这种设计基于一个观察：计算密集型算子具有更复杂的调度，并且难以在短时间内自动生成一个性能超过供应商提供实现的核函数。在极少数情况下，如果一个计算密集型原语图无法匹配供应商库中定义的参数，Korch会拒绝该候选核函数。我们的实现使用cuDNN【【9，cudnn: Efficient primitives for deep learning，2014】】、cuBLAS【【10，Dense Linear Algebra on GPUs，2016】】和TensorRT【【32，NVIDIA TensorRT: Programmable inference accelerator，2017】】作为计算密集型算子的潜在后端。

BLP求解器
Korch通过构建§4中确定的约束，将核函数编排形式化为一个二元线性规划（BLP）问题，并使用PuLP【【21，Pulp: a linear programming toolkit for python，2011】】，一个Python线性规划求解器来解决该BLP问题。在我们测试用例中，最大的子图有584个执行状态和3078个候选核函数。对于我们评估中的所有原语图，PuLP都能在1000秒内收敛到最优解。

5.3 可执行文件生成器
核函数拼接
在生成了最优的核函数编排策略和所有选定候选核函数的CUDA代码之后，Korch根据这些核函数之间的数据依赖关系将它们的实现拼接在一起。因此，整个计算图的端到端执行延迟与ILP任务的目标函数（即公式(2)）是一致的。请注意，Korch只考虑编排后核函数的顺序执行，并未考虑诸如CUDA多流之类的核间优化。通过修改ILP问题公式，可以考虑更高级的核函数执行策略，我们将其作为未来的工作。

A4 实验

实验环境

实验结果

端到端性能比较 (6.2节)
* 实验内容: 在单GPU上评估端到端推理延迟，并将Korch与PyTorch 2.0.1、TVM 0.11和TensorRT 8.2进行比较。
* 实验结果:
* Korch在V100上实现了高达1.7倍的加速，在A100上高达1.6倍。平均加速比分别为1.39倍（V100）和1.30倍（A100）。
* 对于Vision Transformer模型（Segformer, EfficientViT），Korch仍能将延迟分别降低高达1.7倍和1.4倍。
* 对于传统CNN模型（Candy, YOLOv4, YOLOX-Nano），Korch平均性能优于现有框架1.31倍。
* 分析与结论:
* Korch在V100上的性能提升比A100更显著。原因有二：1) Korch不仅优化内存密集型算子，还通过优化数据布局等方式提升计算密集型算子的性能。2) Korch依赖的TVM在A100上的性能不如高度优化的TensorRT，这表明通过利用更先进的核函数生成技术，Korch仍有提升空间。

图 6. 在V100和A100上的端到端性能比较。

算子分裂有效性研究 (6.3节)
* 实验内容: 将Korch的算子分裂后的原语图直接输入TensorRT，以验证算子分裂本身的有效性。实验在Segformer模型上进行。
* 实验结果: 仅应用算子分裂，就使得TensorRT在V100上的性能提升了1.24倍（如图7）。
* 分析与结论: 算子分裂本身就能为现有优化器（如TensorRT）解锁新的优化机会，从而带来性能提升。

图 7. 算子分裂的加速效果：将原语图而不是计算图输入TensorRT带来了1.24倍的加速。

案例研究 (6.4节)
1. 冗余计算 (EfficientViT attention block):
* 策略对比: TensorRT采用直接的核映射策略（图8a）。Korch首先进行图变换（图9）合并ReduceSum和MatMul，然后进行核编排。Korch通过机会性地将Transpose与MatMul融合来优化数据布局（图8b中的k5），使得MatMul性能提升3.52倍。此外，Korch允许Reshape+Transpose+Reshape算子被冗余计算三次（图8b中的k2,k3,k4），从而减少了总核函数数量。
* 结果: Korch仅用7个核函数就完成了计算，比TensorRT的12个核函数少了5个，整个子图实现了3.29倍的加速（图10）。这种优化超出了现有工作的优化空间。

图 8. EfficientViT注意力块上核函数编排策略的比较。在核函数编排之前，Korch首先应用了图9中的图变换。

图 9. 从图8a到图8b的原语图变换。相同颜色的节点表示对原语图的一次变换。ReduceSum和MatMul的合并与图2类似。

图 10. EfficientViT注意力块的案例研究：Korch的策略（图8b）比TensorRT的策略（图8a）实现了3.29倍的加速。

  1. 一个算子映射到不同核函数 (Segformer, Candy):

    • 策略: 对于Segformer，Korch分解Softmax算子，并将其原语映射到四个不同的核函数中，在自注意力模块上比TensorRT快1.50倍。对于Candy，Korch打破了算子边界，分解了InstanceNorm，并将其部分计算与后续的ReLU和Pad算子融合，实现了1.32倍的加速（图12）。
    • 结论: 算子分裂和灵活的核编排允许跨越传统算子边界进行融合，发现更优的执行策略。

      图 12. Candy中一个常见子图模式的核函数编排策略和核函数延迟比较。图12b中的红色框是InstanceNorm的算子分裂。Korch在该子图上比TensorRT快1.32倍。

  2. 贪婪融合可能是次优的 (Segformer):

    • 策略对比: 对于Segformer的一个子图，所有算子都是内存密集型且可融合的。TVM会贪婪地将整个子图融合成一个核函数（图11a）。Korch则会根据情况选择不同策略。
    • 结果: 在batch size为1时，TVM的策略是高效的。但在batch size为16时，TVM生成的融合核函数性能不佳，此时Korch选择使用多个核函数（图11b）的策略，带来了2.24倍的加速（图13）。
    • 结论: 贪婪融合并非总是最优，需要根据具体负载和硬件特性进行全局优化。

      图 11. Segformer子图的两种不同核函数编排策略。TVM总是选择策略(a)，但Korch在batch size为1时选择策略(a)，在batch size为16时选择策略(b)。

      图 13. Segformer子图上不同核函数编排策略的比较。TVM总是选择策略A，但Korch在batch size为1时选择策略A，在batch size为16时选择策略B。

调优时间 (6.5节)
* 结果: 表2显示了各模型的原语图节点数、候选核函数数和总调优时间。实际候选核函数数量远小于理论上的二次方复杂度。最长的调优时间为12.2小时（YOLOv4）。
* 分析: 大部分时间消耗在TVM MetaScheduler对内存密集型核函数的调优上。调优是一次性成本，且可通过数据库缓存和并行化来管理，因此在实际应用中是可接受的。
表 2. 所有评估模型在Perlmutter上的原语图节点数、候选核函数数和端到端调优时间。

A5 相关工作与未来展望

7. 相关工作 (缩写)

8. 未来工作

A6 结论

本文识别了张量程序编译中的核函数编排任务，并提出了Korch，一种发现张量程序最优核函数编排策略的系统性方法。Korch应用算子分裂将张量算子分解为细粒度的原语，将核函数编排形式化为一个二元线性规划任务，并使用现成的求解器生成最优策略。在现代GPU架构上，Korch的性能比现有的张量程序优化器高出1.7倍。

A7 附录

A.1 摘要

此工件附录旨在帮助读者运行工件并复现论文《利用Korch实现张量程序的最优核函数编排》中的结果。

A.2 工件清单（元信息）

A.3 描述

A.4 安装

请参阅README中的“环境准备”部分。

A.5 实验工作流

请参阅README中的“运行TensorRT基线”和“运行Korch”部分。包括以下实验：
* 对TensorRT进行算子分裂的适应性研究。这将花费大约5分钟。
* 在Candy模型上对Korch进行端到端评估。在默认配置下，这将花费大约3小时。
* 在Segformer模型上对Korch进行端到端评估。在默认配置下，这将花费大约7小时。

A.6 实验定制

此工件的默认配置未启用TensorRT后端，因为这会使调优时间至少增加一倍，而速度提升甚微。尽管可以通过在framework/calc.py的第13行将enable_trt更改为True来启用TensorRT，我们建议用户使用默认配置，以节省调优时间，同时不影响Korch系统的功能完整性测试。

A.7 注意事项

Korch的端到端调优时间大部分将消耗在候选核函数分析上。会有一个tqdm进度条显示预计完成时间（ETA）。Korch使用TVM数据库在候选核函数分析期间存储已分析的结果，因此用户可以在此过程中随时暂停并恢复程序。

A.8 评估和预期结果

对于对TensorRT的适应性研究，结果应与图7相似。对于Candy和Segformer的端到端结果，工件默认配置的结果将略慢于图6，因为TensorRT后端被禁用了。

引用文献汇总