DIETCODE: AUTOMATIC OPTIMIZATION FOR DYNAMIC TENSOR PROGRAMS

DIETCODE: 动态张量程序的自动优化

Bojian Zheng 1 2 3, Ziheng Jiang 4, Cody Hao Yu 2, Haichen Shen 5, Josh Fromm 6, Yizhi Liu 2, Yida Wang 2, Luis Ceze 7 6, Tianqi Chen 8 6, Gennady Pekhimenko 1 2 3

A1 主要贡献

本文旨在解决现有自动调度器（auto-scheduler）在处理动态形状（dynamic-shape）工作负载时面临的巨大挑战，即极长的自动调度时间。现有自动调度器由于其搜索空间是依赖于形状的（shape-dependent），因此无法为工作负载所有可能的形状构建一个统一的搜索空间，必须对每个形状单独进行优化，这在实际应用中非常耗时。

为应对此挑战，本文提出了 DietCode，一个新的自动调度器框架，通过构建一个形状通用（shape-generic）的搜索空间和成本模型来高效支持动态形状工作负载。在 DietCode 的设计下，所有形状都在同一个空间内联合搜索并共同更新同一个成本模型，从而显著提高了自动调度的效率。

本文的主要贡献如下：

  1. 识别并解决了关键挑战：发现了将自动调度应用于动态形状工作负载的核心挑战，即形状依赖的搜索空间，并提出使用形状通用的搜索空间来解决此问题。
  2. 构建了 DietCode 框架：这是一个专为动态形状工作负载设计的新型自动调度器框架。它采用联合学习（joint learning）方法，在同一个形状通用的搜索空间内，对工作负载所有可能的形状进行集体优化，并学习同一个形状通用的成本模型。
  3. 显著的性能和效率提升：在对先进的语言模型 BERT 的评估中，DietCode 展示了其巨大优势：
    • 减少自动调度时间：在均匀采样的动态形状上，与最先进的自动调度器 Ansor 相比，DietCode 将自动调度时间减少了高达 5.88倍（如果包含所有可能的形状，预计可达 94.1倍）。
    • 提升运行时性能：与 Ansor 相比，性能提升高达 69.5%；与供应商库（cuBLAS 和 cuDNN）相比，性能提升高达 18.6%。

这些优势使得 DietCode 成为一个高效且实用的动态形状工作负载解决方案。

A3 背景知识与动机

2.1 供应商库和现有自动调度器

高性能算子实现的挑战与方案。为计算密集型算子（如卷积和矩阵乘法）实现高性能一直是一项艰巨的任务。因此，大多数先进的机器学习框架（如 TensorFlow 【1, TensorFlow: A system for large-scale machine learning, 2016, OSDI】, PyTorch 【41, PyTorch: An imperative style, high-performance deep learning library, 2019, NeurIPS】, MXNet 【8, MXNet: A flexible and efficient machine learning library for heterogeneous distributed systems, 2015, CoRR】）都依赖于高度优化的、手工制作的供应商库（例如，英特尔 CPU 上的 oneDNN 【39, oneAPI deep neural network library (oneDNN), 2021】；NVIDIA GPU 上的 cuDNN 【11, cuDNN: Efficient primitives for deep learning, 2014, CoRR】和 cuBLAS 【37, cuBLAS :: CUDA toolkit documentation, 2021】）。尽管这些库性能很高，但其开发需要深厚的软硬件专业知识和巨大的工程投入，导致支持新算子和新硬件的发布周期非常长。

自动调度器框架的出现与工作原理。为了在短时间内为不同硬件提供高性能的张量程序，研究者提出了自动调度器框架（如 Ansor 【56, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI】, TVM 【9, TVM: An automated end-to-end optimizing compiler for deep learning, 2018, OSDI】, Halide auto-scheduler 【2, Learning to optimize Halide with tree search and random programs, 2019, ACM Transactions on Graphics】, Tensor Comprehensions 【50, The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated GPU kernels, automatically, 2020, ACM Transactions on Architecture and Code Optimization】）。自动调度器的输入是张量表达式，包括算子规范和输入形状描述（如图 2(a)）。它通过分析表达式自动构建高性能程序。具体流程为：首先，根据张量表达式构建一个包含大量可能实现（即调度，schedules）的搜索空间（图 1(a) 中的 ❶）；然后，建立一个成本模型来预测每个调度的运行时性能（❷）。在整个自动调度过程中，成本模型会根据真实的硬件测量数据不断更新，并用于指导在搜索空间中的探索（❸）。最终生成的调度可以使算子实现比供应商库的对应版本快高达 3.8 倍。

图1. (a) 现有自动调度器与 (b) DietCode 的代码生成流程对比。现有方法对每个形状单独进行自动调度。DietCode 通过让所有形状在同一个形状通用的搜索空间内联合搜索并更新同一个微内核成本模型来解决该问题。
$$ \begin{array}{ll} \text{(a)} & X:[1024,768], W:[2304,768] \\ Y=X W^{\top} & \\ \text{(b)} & X:[16 \times \mathbf{T}, 768], W:[2304,768], \\ & \mathbf{T} \in[1,128] \end{array} $$

2.2 动态张量程序

现有自动调度器的局限性。尽管实现方式各不相同，但据我们所知，现有的自动调度器（【56, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI】,【9, TVM: An automated end-to-end optimizing compiler for deep learning, 2018, OSDI】,【2, Learning to optimize Halide with tree search and random programs, 2019, ACM Transactions on Graphics】,【50, The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated GPU kernels, automatically, 2020, ACM Transactions on Architecture and Code Optimization】）有一个共同的限制：它们要求工作负载是静态的（即所有形状必须在编译时已知）才能进行分析。这一限制使得它们在许多现实场景中难以甚至不切实际地使用，因为自动调度会耗费大量时间（我们稍后将展示，单个算子可能需要约 42 小时）。这些场景包括：

  1. 神经架构搜索 (NAS)（【59, Neural architecture search with reinforcement learning, 2017, ICLR】）：NAS 的目标是在预定义的超参数（如批量大小、隐藏维度）搜索空间中，为给定的训练数据集找到合适的 DNN。每个超参数组合都会导致一个具有不同形状网络层的架构。为了获得最佳搜索结果，需要探索大量的网络架构（例如，在 【59, Neural architecture search with reinforcement learning, 2017, ICLR】中检查了 12,800 个架构）。

  2. 设计本身即为动态：虽然机器学习从业者通常会为部署而冻结模型，但某些模型，特别是序列学习领域的模型，本质上是动态的。例如，机器翻译（【53, Google’s neural machine translation system: Bridging the gap between human and machine translation, 2016, CoRR】,【51, Attention is all you need, 2017, NeurIPS】）、语音识别（【5, Deep Speech 2 : End-to-end speech recognition in English and Mandarin, 2016, ICML】）和语言建模（【12, BERT: Pre-training of deep bidirectional transformers for language understanding, 2019, NAACL-HLT】）都涉及动态序列长度，这些长度根据输入数据样本而变化。每个序列长度同样对应一个具有不同形状网络层的网络。例如，在先进的语言建模应用 BERT 【12, BERT: Pre-training of deep bidirectional transformers for language understanding, 2019, NAACL-HLT】中，序列长度可以取 1 到 128 之间的任何值。

  3. 不同层位置的形状变化：即使模型是静态的，同一类算子在模型中的不同位置也可能呈现出不同的形状。再以 BERT 模型【12, BERT: Pre-training of deep bidirectional transformers for language understanding, 2019, NAACL-HLT】为例：全连接层的隐藏维度可以取 {768, 2304, 3072} 中的值，具体取决于它在模型中的位置。

动态形状带来的挑战。因此，在这些场景中，同一类型的张量程序会呈现出多种形状。现有自动调度器（【56, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI】, 【9, TVM: An automated end-to-end optimizing compiler for deep learning, 2018, OSDI】, 【2, Learning to optimize Halide with tree search and random programs, 2019, ACM Transactions on Graphics】, 【50, The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated GPU kernels, automatically, 2020, ACM Transactions on Architecture and Code Optimization】）只接受静态形状工作负载的限制，给它们在这些场景中的应用带来了挑战，因为编译时间可能极长（在现代 CPU 上，对单个层类型大致估计分别为 4200、42 和 1 小时，对于整个网络则需要 20 倍的时间【56, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI】）。然而，我们观察到，如果能将这些程序分组为一个单一的动态形状工作负载并进行一次自动调度，编译时间可以显著缩短，最多可达 5.88 倍（如果包含所有可能的形状，预计可达 94.1 倍）。

2.3 为何需要一个新的自动调度器框架？

现有方案的不足。本文提出了一个新颖的自动调度器框架，以高效支持动态形状工作负载。一个合理的问题是，为什么需要这样一个新框架，或者说，为什么现有解决方案不足以应对挑战。具体来说，供应商库（【39, oneAPI deep neural network library (oneDNN), 2021】,【11, cuDNN: Efficient primitives for deep learning, 2014, CoRR】,【37, cuBLAS :: CUDA toolkit documentation, 2021】）或现有自动调度器（【56, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI】,【9, TVM: An automated end-to-end optimizing compiler for deep learning, 2018, OSDI】,【2, Learning to optimize Halide with tree search and random programs, 2019, ACM Transactions on Graphics】,【50, The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated GPU kernels, automatically, 2020, ACM Transactions on Architecture and Code Optimization】）及其扩展是否能胜任？下面我们概述了这些替代方法的一些关键挑战，以突显为何现有解决方案在实践中不尽如人意。我们使用的例子是来自 BERT-base 模型【12, BERT: Pre-training of deep bidirectional transformers for language understanding, 2019, NAACL-HLT】的全连接层 $Y = XW^\top$，在一个常用的批大小为 16 的情况下（见图 2(b)）【56, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI】。张量程序的动态性在于序列长度 $T$，它代表语料库中句子的长度。为便于说明，我们将其范围设为 $[1, 128]$。

供应商库的局限性。供应商库（【39, oneAPI deep neural network library (oneDNN), 2021】, 【37, cuBLAS :: CUDA toolkit documentation, 2021】, 【11, cuDNN: Efficient primitives for deep learning, 2014, CoRR】）是手工制作的，包含多个优化过的内核以覆盖所有可能的形状，工作负载在运行时会根据其形状被动态分派到最合适的调度（由硬编码的启发式规则决定）上执行。当工作负载的形状不能完全匹配调度时，在特定硬件或工作负载上的运行时性能会是次优的（在某些病态情况下性能下降高达 13 倍【3, Bringing TVM into TensorFlow for optimizing neural machine translation on GPU, 2018】）。这种次优性的原因可能是将形状填充到硬编码内核的瓦片大小的倍数，这引入了填充开销并降低了计算效率。先前的工作观察到这种情况表现为冗余计算操作【3, Bringing TVM into TensorFlow for optimizing neural machine translation on GPU, 2018】和延迟阶梯效应【57, Towards latency-aware DNN optimization with GPU runtime analysis and tail effect elimination, 2020, CoRR】。这些低效性推动了扩展供应商库中手工内核的需求，但由于这些内核的复杂性，这是一项非常艰巨的任务。

现有自动调度器的挑战。现有自动调度器（【56, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI】,【9, TVM: An automated end-to-end optimizing compiler for deep learning, 2018, OSDI】,【2, Learning to optimize Halide with tree search and random programs, 2019, ACM Transactions on Graphics】,【50, The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated GPU kernels, automatically, 2020, ACM Transactions on Architecture and Code Optimization】）可以通过为每个可能的形状优化调度来克服上述缺点，但将它们直接应用于动态形状工作负载可能具有挑战性。图 1(a) 展示了采用现有自动调度器处理动态形状工作负载的直接工作流程。由于现有自动调度器一次只能处理一个静态形状，因此每个可能的静态形状都需要单独进行自动调度。这在图 1(a) 中表现为多个自动调度器实例（在 $S_i$ 旁边）。由于在配备 16 个 Intel® Xeon® Platinum 8259CL CPU 的现代机器上【40, Intel Xeon Platinum 8259CL @ 2.50GHz, 2020】，每个实例优化单个形状的调度大约需要 0.33 个 CPU 小时，我们经验性地估计，优化图 2(b) 中指定的单个动态形状全连接层工作负载需要 42 个 CPU 小时。

现有扩展方案的不足。尽管已有基于现有自动调度器的扩展被提出来支持动态形状工作负载（【55, [RFC][AutoTVM] selective tuning, 2019】,【48, Nimble: Efficiently compiling dynamic neural networks for model inference, 2021, MLSys】,【52, [RFC] dynamic shape support - graph dispatching, 2019】），但据我们所知，这些先前的工作都面临着非全自动和/或生成低性能张量程序的挑战：(1) 选择性调优 (Selective tuning) 【55, [RFC][AutoTVM] selective tuning, 2019】通过启发式方法将静态形状工作负载分组到聚类中，并为每个聚类单独优化，这需要人类专业知识，因此不是全自动的。(2) Nimble 【48, Nimble: Efficiently compiling dynamic neural networks for model inference, 2021, MLSys】通过让动态形状工作负载采用一个大尺寸形状进行自动调度，并将其调度普遍应用于所有形状。然而，正如我们在第 5 节的评估中所示，一个在大尺寸形状上最优的调度在其他形状上可能是次优的。(3) 分桶 (Bucketing) 【52, [RFC] dynamic shape support - graph dispatching, 2019】将动态范围划分为小的子范围（例如，图 2(b) 中的 $T \in [1, 128]$ 分为 $T \in [1, 8]$, $[9, 16]$ 等），并仅对每个子范围内的最大值进行调优，同时将其余值填充到最大值。这种方法存在两个低效来源：(i) 它需要在主计算前后分别进行额外的填充和切片操作，这带来了性能和存储开销。(ii) 在填充后的张量上执行计算效率低下，因为许多计算是无用的（例如，在前面的分桶例子中，$T = 8$ 需要的计算量是 $T = 1$ 的 7 倍）。这些弱点阻碍了以实用和高效的方式使用现有自动调度器（【56, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI】,【9, TVM: An automated end-to-end optimizing compiler for deep learning, 2018, OSDI】,【2, Learning to optimize Halide with tree search and random programs, 2019, ACM Transactions on Graphics】,【50, The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated GPU kernels, automatically, 2020, ACM Transactions on Architecture and Code Optimization】）来支持动态形状工作负载，从而催生了新的自动调度器设计。

A3 DIETCODE: 核心思想

本节深入探讨现有自动调度器（【56, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI】,【9, TVM: An automated end-to-end optimizing compiler for deep learning, 2018, OSDI】,【2, Learning to optimize Halide with tree search and random programs, 2019, ACM Transactions on Graphics】,【50, The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated GPU kernels, automatically, 2020, ACM Transactions on Architecture and Code Optimization】）的挑战，以理解为何它们难以被简单改进以高效处理动态形状工作负载，并阐述我们设计背后的核心思想。

3.1 形状通用的搜索空间

图3. 从输入张量表达式(a)转换而来的循环分块调度(b)，其中T是用户提供的循环范围，t是自动调度器可调优的分块大小参数。
图3. 从输入张量表达式(a)转换而来的循环分块调度(b)，其中T是用户提供的循环范围，t是自动调度器可调优的分块大小参数。

构建统一搜索空间的观察。我们观察到，不同形状的搜索空间实际上可以重叠，从而有潜力形成一个形状通用的（shape-generic）搜索空间，而不是使用现有自动调度器（【56, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI】,【9, TVM: An automated end-to-end optimizing compiler for deep learning, 2018, OSDI】,【2, Learning to optimize Halide with tree search and random programs, 2019, ACM Transactions on Graphics】,【50, The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated GPU kernels, automatically, 2020, ACM Transactions on Architecture and Code Optimization】）逐一调整每个可能的形状。以图 3(b) 中由张量表达式(a)转换而来的循环分块（loop tile）调度为例。在调整分块大小 t 时，现有自动调度器仅考虑 T 的所有因子作为候选，这意味着当 $T = 49$ 时，$t \in \{1, 7, 49\}$，当 $T = 50$ 时，$t \in \{1, 2, 5, ..., 50\}$。然而，任何小于 T 的整数实际上都是有效的候选值，它们构成了一个可供所有可能的 T 使用的形状通用搜索空间。我们从这个例子中观察到，搜索空间的构建实质上是在一个为所有可能形状服务的、庞大而统一的搜索空间与解决额外边界检查（在非完美分块情况下需要，例如当 $T = 49$ 时 $t = 10$）的挑战之间做出权衡。如我们将在 4.1 节中展示，通过在代码生成中仔细处理非完美分块，我们可以避免边界检查的开销。因此，非因子分块候选值也能产生与因子候选值相同甚至更好的性能。我们将在第 5 节展示非因子候选值如何积极影响生成的张量程序的性能。

基于微内核的形状通用搜索空间。基于这些优点，我们构建了一个由微内核（micro-kernels）（图 1(b) 中的 ❹）组成的形状通用搜索空间，以高效支持动态形状工作负载。微内核是一个不完整的程序，它执行完整计算的一个分块（tile）。我们使用硬件约束（例如，最大线程数、共享和本地内存量）而非形状信息来确定微内核候选。这些候选作为构建块，通过重复执行来完成一个工作负载实例（定义为动态形状工作负载的一个静态形状实例）。例如，图 4 展示了微内核 dense_128x128，它计算
$Y=XW^{\top}, X: [128, 768], W: [128, 768]$
，如何被用来执行图 2(b) 中 $T = 64$ 的 $Y = XW^\top$ 实例。这是通过沿空间维度将完整程序分解为 $8 \times 18$ 个片段（因为 $16 \cdot T = 8 \times 128, 2304 = 18 \times 128$）并单独执行每个片段来实现的。因为微内核只实现完整计算的一部分，所以每个微内核都可以轻松地移植到不同形状的多个工作负载实例中，因此它是形状通用的。回顾之前的例子，微内核 dense_128x128 不仅可以用于实现 $T = 64$，还可以用于实现 $T = 1, 2, ..., 127, 128$。

图4. 使用微内核 dense_128x128 实现图 2(b) 中 T=64 的 Y=XW⊤。水平和垂直轴表示Y的输出维度。
图5. 使用微内核 dense_128x128 实现图 2(b) 中 T=60 的 Y=XW⊤。最后一列的瓦片被填充（红色显示）以适应微内核。

微内核带来的新挑战。尽管微内核是一个通用的解决方案，但它的使用也带来了两个新的重要挑战：(1) 如何使用微内核组合成高性能的完整程序，特别是在工作负载不能完美适应微内核的情况下（如图 5 所示）；以及 (2) 如何准确高效地预测基于微内核的完整程序的性能。在本文中，我们使用局部填充（local padding）来解决第一个挑战，即在从全局内存获取输入张量时自动填充局部工作空间（详见 4.1 节），并通过构建一个基于微内核的成本模型来解决第二个挑战。

3.2 基于微内核的成本模型

现有成本模型的局限性。为了高效地搜索高性能微内核，我们使用成本模型来指导搜索过程。现有的自动调度器（【56, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI】, 【9, TVM: An automated end-to-end optimizing compiler for deep learning, 2018, OSDI】, 【2, Learning to optimize Halide with tree search and random programs, 2019, ACM Transactions on Graphics】, 【50, The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated GPU kernels, automatically, 2020, ACM Transactions on Architecture and Code Optimization】）也有成本模型，通过提取程序的特征（例如循环结构、内存访问模式）来预测其计算吞吐量，但这些成本模型只能接受完整的程序作为输入：

$$ \text{Cost}(P) = f(\text{FeatureExtractor}(P)) $$

其中 f 是成本函数（例如 XGBoost 【7, XGBoost: A scalable tree boosting system, 2016, KDD】），P 是完整程序，Cost 是计算吞吐量。然而，由于微内核是不完整的（它们是完整程序的一个分块），它们不能直接应用于现有的成本模型。为了应对这一挑战，我们构建了一个基于微内核的成本模型。

分解成本模型。关键思想是，一个由微内核 M 构成的完整程序 P 的成本可以分解为两部分：(1) 一个形状通用的成本函数 $f_{MK}$，用于预测 M 的成本；(2) 一个形状依赖的适应成本函数 $f_{adapt}$，用于定义将 M 移植到 P 的惩罚。其中 $f_{MK}$ 是一个在自动调度过程中必须通过真实硬件测量来学习和更新的函数，而 $f_{adapt}$ 是一个可以利用核心占用率和填充率来评估的简单项（换句话说，它不需要特征提取，详见 4.2 节）。下面是我们成本模型的数学形式：

$Cost_M (P) = f_{MK}(FeatureExtractor(M)) \cdot f_{adapt}(P, M)$

新成本模型的优势。与公式 1 相比，公式 2 的关键优势在于效率：要预测共享同一微内核 M 的多个完整程序的性能，程序特征提取只需对微内核 M 进行一次，而每个程序只需单独更新其适应成本。这极大地消除了特征提取中的冗余，从而实现了更高效的自动调度流程。

...
表1. 三种方案在自动调优能力和调优运行时复杂度方面的比较。|S| 表示可能形状的数量。

3.3 DietCode 的联合学习

DietCode 框架与优化流程。基于上述的搜索空间和成本模型，我们提出了 DietCode，一个新的自动调度器框架，包含三个关键组件：(1) 一个形状通用的搜索空间，(2) 一个基于微内核的成本模型，以及 (3) 一个分派器（dispatcher）（图 1(b) 中的 ❼，详见 4.3 节），它在运行时根据公式 2 自动将静态形状分派给微内核。DietCode 的优化工作流是一个联合学习过程（图 1(b) 中的 ❻），其中一个动态形状工作负载的所有工作负载实例共享相同的搜索空间，并共同学习同一个成本模型。这个工作流使得 DietCode 能够以类别为基础进行操作，每个类别共享同一个形状通用的程序（即微内核）。

方法对比。表 1 对我们目前描述的三种主要方法进行了比较：供应商库（【39, oneAPI deep neural network library (oneDNN), 2021】, 【37, cuBLAS :: CUDA toolkit documentation, 2021】, 【11, cuDNN: Efficient primitives for deep learning, 2014, CoRR】）、现有自动调度器（【56, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI】, 【9, TVM: An automated end-to-end optimizing compiler for deep learning, 2018, OSDI】, 【2, Learning to optimize Halide with tree search and random programs, 2019, ACM Transactions on Graphics】, 【50, The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated GPU kernels, automatically, 2020, ACM Transactions on Architecture and Code Optimization】）和 DietCode，其中我们用 |S| 表示动态形状工作负载可能形状的数量。由于 DietCode 可以在一次运行中生成所有形状通用的类别，它能够显著减少在动态形状工作负载上与现有自动调度器相比的自动调度时间，我们的评估将在第 5 节中展示这一点。

A2 实现细节

我们将 DietCode 作为 TVM 【9, TVM: An automated end-to-end optimizing compiler for deep learning, 2018, OSDI】 的一个自动调度器子模块进行集成。本节将重点介绍 DietCode 设计的一些细节。

4.1 局部填充

问题：边界检查开销。DietCode 的核心思想之一是将微内核通用地应用于所有工作负载实例，每个实例对应动态形状工作负载的一个静态形状。然而，工作负载实例通常无法完美适配微内核，如图 5 所示，最后一列的微内核未被完全物化。在这些情况下，需要在微内核内部注入边界检查，以确保程序不会操作无效数据，但这也会导致严重的性能下降（在我们的评估中，使用图 5 所示的例子，在现代 Tesla T4 GPU 【34, NVIDIA T4 70W low profile PCIe GPU accelerator, 2020】 上的性能下降高达 17 倍）。这是因为这些检查在生成的程序中引入了额外的分支和计算指令。

图6. (a) 带有边界检查的分块循环。(b-d) 三种优化策略：(b) 全局填充，(c) 循环划分，(d) 局部填充。
图6. (a) 带有边界检查的分块循环。(b-d) 三种优化策略：(b) 全局填充，(c) 循环划分，(d) 局部填充。

解决方案对比。据我们所知，有三种解决方案可以缓解这个问题。为展示它们，我们以图 6(a) 中带有边界检查的代码片段为例，其中我们极大地简化了调度，只保留其骨架以便于理解（从全局内存取输入张量到局部工作空间，计算，并将输出张量写回全局内存）。
1. 全局填充 (Global padding)：在主计算之前填充输入张量，并在之后切片输出张量。这可以消除计算内核中的所有边界检查，但需要额外的存储空间并注入填充/切片操作（见图 6(b)）。
2. 循环划分 (Loop partitioning) (【48, Nimble: Efficiently compiling dynamic neural networks for model inference, 2021, MLSys】)：将完整程序划分为可以消除边界检查的区域和不能消除的区域（见图 6(c)）。循环划分的性能优势取决于内外循环的相对比例：如果 $t \ll T$，则会有更多没有边界条件的 i.0 迭代，从而带来更多好处。然而，如果 $t \ll T$ 不成立，则优势可能有限。
3. 局部填充 (Local padding)：在从全局内存获取输入张量值时填充局部工作空间，并在写回时切片局部工作空间。这相当于在获取和写回阶段保留边界检查，而在计算阶段移除它们（见图 6(d)）。局部填充的思想基于两个关键观察：(i) 只有计算阶段的边界检查对运行时性能有主导影响，因为获取和写回阶段的检查可以被全局内存传输的延迟所掩盖【33, CUDA fundamental optimization, part 1, 2019】。(ii) 在计算阶段计算出的填充数据值将被写回阶段的边界检查过滤掉，因此它们不影响输出张量的值（即程序正确性）。

选择局部填充的原因。在本文中，我们选择了第三种方案，因为它兼具透明（不产生存储开销和额外算子）和可通用地应用于大小形状（而循环划分仅在大形状上效果最佳）的优点。我们在第 5 节的评估中证明，局部填充是解决边界检查挑战的有效方案。

4.2 基于微内核的成本模型

成本模型的设计。为了准确预测基于微内核的完整程序的性能，我们在成本模型中设计了三个项，分别用于解释：① 微内核的性能，② 硬件核心占用率惩罚（与执行该微内核以构成完整程序的次数相关），以及 ③ 填充惩罚。成本模型的数学表达式如下：

$Cost_M (P) = f_{MK}(FeatureExtractor(M))① \cdot$ $$\textcircled{2}f_{\text{occ}}(P/M) \cdot f_{\text{pad}}(P, M)\textcircled{3} \atop \underbrace{\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad}_{f_{\text{adapt}}(P,M)}$$

其中 M 表示微内核，P 表示完整程序。适应成本 $f_{adapt}$ 的两个函数 $f_{OCC}$ 和 $f_{padding}$ 分别对应：

硬件核心占用率惩罚 $f_{OCC}$。我们使用线性回归模型来建模此项：

$$f_{\text{occ}}(P/M) = k \frac{P/M}{\text{ceil\_by}(P/M, \text{NumCores})} + b$$

其中，项 P/M 对应于执行 M 形成 P 的次数，系数 k 和 b 是可学习的参数。由于每个微内核被分派到一个硬件核心，P/M 项也与被占用的硬件核心数量相关。因此，上式中的项 ④ 表示由 M 构成的完整程序 P 的硬件占用率。该比率再由两个系数 k 和 b 加权，其值在自动调度时动态适应硬件平台，以估计硬件占用成本。实际上，由于我们希望当所有核心都被完全占用时占用成本为 1，我们可以进一步推导：

$$f_{\text{occ}}(P,M) = 1, \text{ when } \frac{P/M}{\text{ceil\_by}(P/M, \text{NumCores})} = 1$$
$$\Rightarrow k+b=1 \Rightarrow b=1-k$$

以减少一个动态参数 b。

填充惩罚 $f_{pad}$。此项可以以硬件无关的方式建模为 pad_by(P,M)/P（pad_by 原语将完整程序 P 填充到 M 的大小）。

模型与实现的对应关系。图 7 说明了公式 3 中三个项与微内核如何构成完整程序之间的相关性，我们可以看到公式与完整程序构成之间的一一对应关系。这种对应关系使我们能够准确预测基于微内核的完整程序的性能，从而在联合学习过程（图 1 中的 ❻）中使用进化搜索（evolutionary search）【52, Evolutionary algorithms: A critical review and its future prospects, 2016, ICGTSPICC】高效地搜索高性能微内核。

图7. fMK, fOCC, 和 fpad 解释了不同的程序行为（示例取自图5）。
图7. fMK, fOCC, 和 fpad 解释了不同的程序行为（示例取自图5）。

4.3 自动分派

分派机制。联合学习过程完成后，DietCode 会生成一组微内核作为自动调度的结果。为了将所有可能的形状分派到这些微内核，我们让每个形状 S 根据公式 3 中的成本公式为其“最喜欢”的微内核投票：

$\text{vote}(S) = \operatorname{argmax}_M (\text{Cost}_M (P(S, M)))$

其中，LHS 表示投票选出的微内核，项 $P(S, M)$ 指的是使用微内核 M 构成的形状为 S 的完整程序。在所有 S 都投票后，我们使用 scikit-learn 框架【42, Scikit-learn: Machine learning in Python, 2011, JMLR】训练一个决策树（图 1 中的 ❼）。决策树的输入是所有可能的形状，输出标签是它们选择的微内核。生成的决策树会自动将选择相同微内核的形状归为一类，并以 C 风格源代码格式导出，以便在运行时根据输入形状高效地分派到其对应的微内核。

现在我们已经介绍了 DietCode 设计的细节，接下来将在最先进的机器学习工作负载上评估其效率。

A4 实验环境

  • 硬件配置：实验平台为 Amazon EC2 G4dn 实例【4, Amazon EC2 G4 instances, 2021】，配备 16 核 Intel® Xeon® Platinum 8259CL CPU【40, Intel Xeon Platinum 8259CL @ 2.50GHz, 2020】 和 1 块 NVIDIA Tesla T4 GPU【34, NVIDIA T4 70W low profile PCIe GPU accelerator, 2020】。
  • 软件配置：
    • CUDA 版本：11.3【35, Programming guide :: CUDA toolkit documentation, 2021】
    • cuDNN 版本：8.2【38, Developer guide :: NVIDIA deep learning cuDNN documentation, 2021】
    • 基础框架：TVM v0.8.dev0【9, TVM: An automated end-to-end optimizing compiler for deep learning, 2018, OSDI】
  • 模型与数据集：
    • BERT-base【12, BERT: Pre-training of deep bidirectional transformers for language understanding, 2019, NAACL-HLT】：用于端到端模型评估，重点关注动态序列长度。
    • Dense 和 Batched Matrix Multiplication (BatchMatmul) 层：从 BERT-base 中提取的算子，同样使用动态序列长度进行性能分析。
  • 实验设置：
    • 由于对所有可能的动态形状（范围 [1, 128]）进行详尽的自动调度不切实际（单个工作负载可能耗时 42 CPU 小时），实验在 [1, 128] 范围内均匀采样了 8 个形状配置进行性能和调度时间的比较，并基于此推算整个范围的总调度时间。
  • 基线方法：
    • Vendor：NVIDIA GPU 上的供应商库，即 cuBLAS【37, cuBLAS :: CUDA toolkit documentation, 2021】 和 cuDNN【11, cuDNN: Efficient primitives for deep learning, 2014, CoRR】。
    • Ansor【56, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI】：针对静态形状工作负载的最先进的自动调度器实现。每个静态形状工作负载都经过至少 1000 次试验的充分调优。
    • Nimble【48, Nimble: Efficiently compiling dynamic neural networks for model inference, 2021, MLSys】：针对动态代码生成的最先进工作，它通过调优最大形状并使用循环划分将其调度通用地应用于所有形状，是 Ansor 的一个扩展。
  • 评估指标：
    • 端到端延迟：整个模型的延迟，单位为毫秒 (ms)，5 次运行的平均值。
    • 算子运行时成本：单个算子的运行时成本，单位为微秒 (µs)，100 次运行的平均值。
    • 总自动调度时间：总耗时，单位为小时 (hours)。

A4 实验结果

5.2 端到端模型评估

性能提升：在具有动态序列长度的 BERT-base 模型上，DietCode 的端到端延迟表现出色。如图 8 所示，与 Ansor 相比，DietCode 的延迟最多降低了 69.5%（平均降低 29.9%），与 Vendor 库相比，最多降低了 18.6%（平均降低 5.4%）。性能提升的原因将在后续对单个层的分析中揭示。

图8. Vendor、Ansor 和 DietCode 在具有动态序列长度的整个 BERT-base 模型上的延迟比较。
图8. Vendor、Ansor 和 DietCode 在具有动态序列长度的整个 BERT-base 模型上的延迟比较。

调度时间缩短：图 9 展示了 Ansor 和 DietCode 在采样序列长度上的自动调度时间。结果显示，对于整个 BERT 模型，DietCode 的调度时间比 Ansor 在采样形状上快 5.88 倍。由于 Ansor 需要为每个形状单独调度，而 DietCode 只需调度一次，因此推算到整个 [1, 128] 范围，这一速度优势预计将扩大到 94.1 倍。

图9. Ansor 和 DietCode 在各种动态工作负载上的自动调度时间比较。越低越好。
图9. Ansor 和 DietCode 在各种动态工作负载上的自动调度时间比较。越低越好。

5.3 动态全连接层 (Dense Layer)

运行时性能：在图 2(b) 所示的动态序列长度的全连接层上，DietCode 的性能同样优于所有基线。如图 10 所示，DietCode 的平均运行时间比 Ansor 少 30.5%，比 Nimble 少 23.9%，比 Vendor 少 5.3%。DietCode 相对于 Ansor 和 Nimble 的性能优势源于其形状通用的搜索空间，该空间能够探索到被 Ansor 形状依赖的构造所忽略的程序。由于 cuBLAS 是闭源库，无法分析 DietCode 优于 Vendor 的确切原因。

Dense([16 × T, 768], [2304, 768]) 图10. Vendor、Ansor、Nimble 和 DietCode 在具有动态序列长度的 Dense 层上的运行时比较。
Dense([16 × T, 768], [2304, 768]) 图10. Vendor、Ansor、Nimble 和 DietCode 在具有动态序列长度的 Dense 层上的运行时比较。

调度时间：如图 9 所示，对于动态全连接层，DietCode 在采样序列长度上的自动调度时间比 Ansor 少 6.3 倍（预计在整个范围上少 100.8 倍）。虽然低于理论上的 8 倍，这可能是因为形状通用的搜索空间构建和探索比形状依赖的更复杂（需要检查硬件约束），但调度时间的改进仍然非常显著，并且随着动态序列长度 T 的取值范围扩大，这一优势会更加明显。

$$ Y = \text{BatchMatmul}^{\text{NT}}(X, W) $$
$$ X : [192, \mathbf{T}, 64], W : [192, \mathbf{T}, 64], \mathbf{T} \in [1, 128] $$
$$ Y = \text{BatchMatmul}^{\text{NN}}(X, W) $$
$$ X : [192, \mathbf{T}, \mathbf{T}], W : [192, \mathbf{T}, 64], \mathbf{T} \in [1, 128] $$

5.4 动态批处理矩阵乘法层 (BatchMatmul Layer)

运行时性能：实验在图 11 定义的动态批处理矩阵乘法工作负载上进行，该负载由于有多个动态轴而更具挑战性。实验评估了 NT 和 NN 两种数据布局，以验证 DietCode 对动态空间轴和归约轴混合情况的通用性。
如图 12 所示：
* 在 NT 布局下，DietCode 的平均运行时间比 Ansor 好 21.5%，比 Nimble 好 351%，比 Vendor 好 12.3%。
* 在 NN 布局下，DietCode 的平均运行时间比 Ansor 好 24.2%，比 Nimble 好 40.4%，比 Vendor 好 15.4%。

对 Nimble 性能的分析：Nimble 在 NT 布局下性能不佳的原因有两点：(1) 它直接将最大形状上生成的调度应用于其他形状，由于填充等问题导致次优性能。(2) 它未能有效处理张量程序中的边界检查，这在该场景下由于有两个动态轴而更为关键。由于形状相对较小，Nimble 使用的循环划分技术无法发挥优势（如 4.1 节所述，没有足够的循环可以被无边界检查地划分）。

图12. Vendor、Ansor、Nimble 和 DietCode 在具有动态序列长度的 BatchMatmul 上的运行时比较。
图12. Vendor、Ansor、Nimble 和 DietCode 在具有动态序列长度的 BatchMatmul 上的运行时比较。

调度时间：如图 9 所示，DietCode 的自动调度时间在 NT 和 NN 布局上分别比 Ansor 少 4.55 倍和 5.56 倍。这种效率使得在合理时间内完成整个模型的端到端调优成为可能。

总结：以上实验证明，DietCode 在保持生成高性能张量程序能力的同时，极大地提高了自动调度的效率。它能通用地应用于多动态轴、混合动态轴以及整个模型，使其成为动态形状工作负载的实用解决方案。

A7 相关工作

与现有自动调度器的关系。DietCode 通过构建一个形状通用的搜索空间，并让所有可能的形状在同一空间内联合搜索、更新同一个成本模型，解决了自动调度动态形状工作负载的关键挑战。这一点在现有的、针对静态形状工作负载的自动调度器中是未见的（例如 Ansor 【56, Ansor: Generating high-performance tensor programs for deep learning, 2020, OSDI】, TVM 【9, TVM: An automated end-to-end optimizing compiler for deep learning, 2018, OSDI】, Halide auto-scheduler 【2, Learning to optimize Halide with tree search and random programs, 2019, ACM Transactions on Graphics】, TensorComprehensions 【50, The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated GPU kernels, automatically, 2020, ACM Transactions on Architecture and Code Optimization】）。尽管 DietCode 目前是在 TVM 【9, TVM: An automated end-to-end optimizing compiler for deep learning, 2018, OSDI】之上实现的，但其思想可以普遍应用于其他编译器框架，如 Halide 【43, Halide: Decoupling algorithms from schedules for high-performance image processing, 2018, Communications of the ACM】, TensorComprehensions 【50, The next 700 accelerated layers: From mathematical expressions of network computation graphs to accelerated GPU kernels, automatically, 2020, ACM Transactions on Architecture and Code Optimization】, Tiramisu 【6, Tiramisu: A polyhedral compiler for expressing fast and portable code, 2019, CGO】, XLA 【45, XLA: Compiling machine learning for peak performance, 2020】, MLIR 【29, MLIR: Scaling compiler infrastructure for domain specific computation, 2021】, Glow 【44, Glow: Graph lowering compiler techniques for neural networks, 2018, CoRR】, taco 【26, taco: A tool to generate tensor algebra kernels, 2017, ASE】, 和 TASO 【25, TASO: optimizing deep learning computation with automatic generation of graph substitutions, 2019, SOSP】。

基于重用的调优器。Selective Tuning 【55, [RFC][AutoTVM] selective tuning, 2019】 和 ETO 【13, ETO: Accelerating optimization of DNN operators by high-performance tensor program reuse, 2021, VLDB Endowment】 根据一组预定义规则（例如 Selective Tuning 中的相似度比率）将工作负载分组到聚类中，并在单个聚类内重用相同的调度。它们的方法与 DietCode 采用的联合学习是并行（parallel）的。

自动调优器。AutoTVM 【10, Learning to optimize tensor programs, 2018, NeurIPS】, ProTuner 【23, ProTuner: Tuning programs with monte carlo tree search, 2020】, 和 FlexTensor 【57, FlexTensor: An automatic schedule exploration and optimization framework for tensor computation on heterogeneous system, 2020, ASPLOS】 利用程序模板来定义搜索空间并指导搜索过程。它们与自动调度器的不同之处在于搜索空间必须手动定义。然而，如果存在一个预定义的微内核搜索空间，DietCode 的核心思想也可以应用于这些自动调优器。

动态神经网络。动态批处理（Dynamic batching）是 DyNet 【32, DyNet: The dynamic neural network toolkit, 2017】, Cavs 【54, Cavs: An efficient runtime system for dynamic neural networks, 2018, USENIX ATC】, BatchMaker 【22, Low latency RNN inference with cellular batching, 2018, EuroSys】, 和 TensorFlow Fold 【30, Deep learning with dynamic computation graphs, 2017, ICLR】 等框架在批处理大小动态变化时常用的一种图级优化。Nimble 【48, Nimble: Efficiently compiling dynamic neural networks for model inference, 2021, MLSys】 和 DISC 【58, DISC: A dynamic shape compiler for machine learning workloads, 2021, CoRR】 都设计了编译器来表示和执行动态神经网络。Cortex 【21, Cortex: A compiler for recursive deep learning models, 2020, CoRR】 是一个基于编译器的递归神经网络框架。这些工作专注于图级优化，因此与在单个层级操作的 DietCode 是正交的。实际上，这些图级解决方案也可以利用 DietCode 来进行高效的算子代码生成。

A5 结论

本文提出了 DietCode，一个专为动态形状工作负载设计的新型自动调度器框架。我们的评估表明，在一个完整的、最先进的 DNN 模型端到端测试中，DietCode 能够将自动调度时间显著减少 5.88 倍（如果包含所有可能的形状，预计可减少 94.1 倍），同时性能比最先进的自动调度器高出 69.5%，比供应商库高出 18.6%。我们希望 DietCode 能成为一个高效的平台，为关键机器学习应用的高效系统设计提供进一步的研究基础。

A6 附录

A.1 摘要

我们提供了与第 3 节和第 4 节对应的源代码和脚本作为我们的工件（artifact）。我们使用 AWS G4 实例作为硬件平台，并使用 NVIDIA 驱动和 Docker 来构建软件栈。安装后，每个实验都由一个单独的脚本文件自动化执行。

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

  • 算法: DietCode 自动调度工作流
  • 程序: TVM 【9, TVM: An automated end-to-end optimizing compiler for deep learning, 2018, OSDI】 子模块（一个集成了 DietCode，另一个有少量改动）+ 基准测试用例
  • 编译: 请使用仓库中提供的脚本文件 scripts/1-compile.sh 进行编译。
  • 转换: 不适用
  • 二进制文件: 不适用
  • 数据集: 不适用
  • 运行时环境: 请使用提供的 Dockerfile 在 AWS 深度学习 AMI (Ubuntu 18.04) 版本 50.0 下构建一个容器化环境。
  • 硬件: Amazon EC2 G4 实例 (g4dn.4xlarge)
  • 运行时状态: 硬件资源（CPU、GPU、RAM、PCIe）无其他进程争用。
  • 执行: 请使用仓库中为每个实验提供的脚本文件 scripts/2-experiment_*.sh。
  • 指标: 计算吞吐量 (TFLOPs/s) 和墙上时钟时间
  • 输出: CSV 文件
  • 实验: 对具有动态序列长度的 Dense 和 BatchMatmul 算子进行自动调度。对 BERT 【12, BERT: Pre-training of deep bidirectional transformers for language understanding, 2019, NAACL-HLT】 模型的关键算子进行自动调度（同样具有动态序列长度）。
  • 所需磁盘空间 (大约): 至少 20 GB（Docker 镜像 15 GB，编译产物 1.5 GB）
  • 准备工作流所需时间 (大约): 1 小时
  • 完成实验所需时间 (大约): 完整实验需要 10 小时
  • 是否公开可用: 是
  • 代码许可证 (如果公开): 不适用
  • 数据许可证 (如果公开): 不适用
  • 使用的工作流框架: 不适用
  • 已归档 (提供 DOI): https://doi.org/10.5281/zenodo.6326726

A.3 描述

A.3.1 如何交付

源代码已在 GitHub 上公开：
https://github.com/UofT-EcoSystem/DietCode
分支：MLSys2022_AE

A.3.2 硬件依赖

我们使用 AWS G4 实例（类型为 g4dn.4xlarge）来评估我们的结果。

A.3.3 软件依赖

主要的软件依赖是 NVIDIA GPU 驱动和 NVIDIA Docker 容器工具包（nvidia-docker2 和 nvidia-container-runtime）。这两者都已自动包含在 AWS 深度学习 AMI (ubuntu 18.04) 版本 50.0 中。其余依赖项可以通过构建 Docker 镜像获得。为了方便 Docker 命令的使用，我们使用了 docker-compose，它是 Docker 的一个包装器。

A.3.4 数据集

（原文无内容）

A.4 安装

请按照以下步骤操作：

  • 克隆项目：git clone https://github.com/UofT-EcoSystem/DietCode -b MLSys2022_AE
  • 安装 docker-compose：sudo -H pip3 install docker-compose
  • 构建包含运行实验所需所有软件依赖的 Docker 镜像：DietCode$ docker-compose build tvm-dev` * 根据镜像创建一个运行中的容器：`DietCode$ docker-compose run --rm tvm-dev
  • 构建 DietCode 和 TVM 基线：
    /mnt$ ./scripts/1-compile.sh tvm` `/mnt$ ./scripts/1-compile.sh tvm_base

A.5 实验工作流

  • 具有动态序列长度的 Dense 层 (正文 5.3 节)
    /mnt$ ./scripts/2_1-experiment_dynamic_dense.sh` * **具有动态序列长度的 BatchMatmul 层 (正文 5.4 节)** `/mnt$ ./scripts/2_2-experiment_dynamic_batch_matmul_nt.sh
    /mnt$ ./scripts/2_3-experiment_dynamic_batch_matmul_nn.sh` * **具有不同序列长度的 BERT (正文 5.2 节)** `/mnt$ ./scripts/2_4-experiment_bert.sh

A.6 评估与预期结果

每个实验运行后，将在各自的文件夹 ops/dense, ops/batch_matmul 和 networks/bert 中生成一个名为 temp_workspace.csv 的 CSV 文件，报告延迟数据（单位为秒，越低越好）。同时，在同一文件夹中会生成 ansor/dietcode_autosched_timer.csv，报告完成自动调度过程所需的时间（同样单位为秒，越低越好）。

A.7 实验定制

不适用

A.8 注意事项

请注意，整个自动调度工作流需要时间来完成。因此，可以使用 AUTO_SCHED_NTRIALS=200 ./scripts/... 前缀来使用较少次数的自动调度试验。生成的张量程序在功能上仍然是正确的，但性能可能是次优的。

A.9 方法论

提交、评审和徽章方法论：
* http://cTuning.org/ae/submission-20190109.html
* http://cTuning.org/ae/reviewing-20190109.html
* https://www.acm.org/publications/policies/artifact-review-badging