Optimizing Dynamic-Shape Neural Networks on Accelerators via On-the-Fly Micro-Kernel Polymerization

作者/机构: Feng Yu, Guangli Li, Jiacheng Zhao, Huimin Cui, Xiaobing Feng, Jingling Xue (SKLP, ICT, CAS, UCAS, UNSW)

A1 主要贡献

本文旨在解决动态形状神经网络中张量程序优化的挑战。现有方法存在编译成本高、无法适应任意形状变化等问题。为实现对动态形状神经网络的高效优化，本文提出了MikPoly，一个基于微内核聚合（micro-kernel polymerization）的新型动态形状张量编译器。

核心问题与研究目标：
- 问题: 传统的静态形状张量编译器（如TVM）虽然能为固定形状生成高性能代码，但其搜索过程耗时过长，不适用于算子形状在运行时才确定的动态场景。现有的动态形状编译器（如DietCode）虽然加快了编译速度，但要求预先定义形状的变化范围，这限制了其在形状变化频繁且任意的场景中的应用。同时，厂商提供的手工优化库（如cuBLAS）在不同张量形状上性能波动巨大，难以保证持续的最优性能（如图1所示）。
- 目标: 设计一个能够为任意运行时确定的张量形状，高效生成优化张量程序的编译器框架，以提升动态形状神经网络在现代加速器上的执行性能。


图 1. GEMM在NVIDIA A100 GPU上不同形状的性能表现（使用cuBLAS）。

创新点与主要贡献：
1. 提出了一种两阶段优化方法：该方法将底层优化问题解耦。
* 离线阶段：为一组固定的形状创建高度优化的微内核（micro-kernels）。
* 在线阶段：当运行时形状已知时，通过“聚合”（polymerize）多个微内核来为任意形状生成优化的程序。
2. 引入了一个精确且轻量级的代价模型：该模型用于指导高效的在线聚合。
* 离线阶段：通过同时考虑计算和内存访问行为，对单个微内核的性能进行建模。
* 在线阶段：在考虑并行性的基础上，评估不同聚合策略和模式下各种程序实现的性能，从而选择最优方案。
3. 实现了MikPoly编译器并进行了全面评估：在GPU张量核（Tensor Cores）和昇腾NPU两种代表性加速器上，对MikPoly进行了实现和评估。实验结果表明，MikPoly在GEMM和卷积算子上，相比于顶尖的厂商库，在GPU和NPU上分别取得了平均1.29倍（峰值11.05倍）和1.70倍（峰值15.32倍）的性能提升。

A3 背景知识与设计原则

本节首先解释了优化动态形状算子的重要性，然后回顾了现有解决方案，并以通用矩阵乘法（GEMM）为例介绍了本文的解决思路。

2.1 动态形状神经网络

传统的深度神经网络（DNN）通常使用静态模型结构，即每个算子的输入输出张量形状是固定的。然而，现实世界的应用常常表现出动态行为，例如语言模型中长度可变的句子，这使得静态形状神经网络不足以应对。动态形状神经网络被提出来支持更复杂的智能应用，其代表性场景包括：

  1. 动态批量大小（Dynamic Batch Sizes）：在模型训练中，批量大小是一个关键参数。动态调整批量大小可以在训练过程中更好地权衡收敛速度、稳定性与计算资源消耗，从而优化训练过程。
  2. 动态图像分辨率（Dynamic Image Resolution）：在计算机视觉任务中，图像由于分辨率不同而具有变化的张量形状。虽然可以将图像缩放到固定尺寸，但这会牺牲原始信息。像Faster R-CNN 【索引17，Fast r-cnn，2015，IEEE international conference on computer vision】这样的先进模型使用支持动态形状输入的池化方法，以有效处理不同形状的图像，实现对小物体的精确检测。
  3. 动态序列长度（Dynamic Sequence Length）：在自然语言处理应用中（如BERT 【索引10，Bert: Pre-training of deep bidirectional transformers for language understanding，2018，arXiv】），输入句子长度的变化导致张量形状动态改变。一种常见的解决方案是将所有序列填充到预定义的统一最大长度，但这可能导致当序列远小于最大长度时造成资源浪费。

2.2 现有技术水平

自动调度器（如TVM 【索引7，TVM: An automated end-to-end optimizing compiler for deep learning，2018，13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)】）通过探索特定形状的搜索空间来生成高性能的张量程序。图2展示了现有静态和动态形状张量编译器优化GEMM的过程。


图 2. 现有静态和动态形状张量编译器优化GEMM张量程序。

静态形状张量编译器的工作方式。以一个固定形状(M, N, K) = (4096, 1024, 4096)的GEMM算子（图2中❶）为例，其初始的朴素张量程序是一个三层嵌套循环（❷）。像TVM这样的编译器提供一个分块（tiled）程序模板（❸），其中包含未确定的分块参数（如TM.0, TN.0, TK.0）。编译器通过自动调度过程，在巨大的搜索空间中探索最优的分块大小（❹），最终为该特定形状生成一个精细调优的张量程序（❺）。然而，这个过程非常耗时（例如0.33个CPU小时），只适用于离线编译的静态场景，不适合需要在模型执行期间进行在线编译的动态形状场景。

动态形状张量编译器的工作方式。对于一个动态形状的GEMM（图2中❶），例如(M, N, K) = (τ, 1024, 4096)，其中τ是一个动态维度，其范围由开发者指定（如[1, 4096]）。为避免高昂的编译成本，像DietCode 【索引67，Dietcode: Automatic optimization for dynamic tensor programs，2022，Proceedings of Machine Learning and Systems】这样的编译器会为一组代表性形状生成一系列调优的张量程序（❼），每个程序都针对一组形状而非单一形状进行优化。在运行时，根据已知的张量形状选择一个合适的预编译程序。然而，这种方法要求预先知道张量形状的范围（例如τ ∈ [1, 4096]），这限制了其适用性。Nimble 【索引51，Nimble: Efficiently compiling dynamic neural networks for model inference，2021，Proceedings of Machine Learning and Systems】也存在类似限制。

现有方法的局限性。现有的静态和动态编译器都为特定的输入范围优化算子，导致对于范围外的形状可能出现性能下降或运行时错误，对于范围内形状的性能也可能是次优的。因此，需要一种有效机制来为任意形状提供高性能的张量程序。

2.3 我们的解决方案

MikPoly的核心思想。MikPoly通过一个两阶段的程序模板来创新动态形状张量算子的编译过程，如图3所示。对于一个编译时形状未知的GEMM算子（❶），我们设计了一个程序模板（❷），该模板包含用于创建微内核模板（❸）的离线循环（蓝色部分），以及围绕这些循环的在线循环（橙色部分）。这种设计允许离线创建不同尺寸的优化微内核。通过灵活地重组在线循环（使用不同的聚合模式和策略），我们可以在运行时动态地生成各种GEMM实现。最后，利用一个精确且轻量级的代价模型（❹和❺），选择与运行时确定的动态形状最匹配、性能最佳的GEMM实现。


图 3. MikPoly通过两阶段动态形状张量编译器为GEMM生成优化的张量程序。

离线阶段。MikPoly利用自动调度器，从微内核模板（❸）中创建一组高度优化的固定大小微内核及其性能模型。这个过程类似于静态形状编译器。

在线阶段。一旦GEMM的运行时形状已知（例如(M, N, K) = (4096, 1024, 4096)），MikPoly会动态地将其程序模板（❷）调整为多种GEMM实现。这涉及到探索不同的聚合模式（❹中的模式I和II），并使用不同的聚合策略，从离线生成的固定大小微内核集合中实例化它们的参数化微内核。最终，通过一个精确且轻量级的代价模型（❺）来确定并选择最优的张量程序。这种方法通过将编译时优化的固定大小微内核与不同的聚合模式和策略相结合，高效地为动态形状张量算子生成程序，从而显著提升动态形状神经网络在新型加速器上的性能。

A2 方法细节

图4展示了MikPoly的概览，它包含两个核心阶段：微内核生成（S1）和微内核聚合（S2）。MikPoly通过一个多级加速器抽象来对目标设备建模，其中每个处理单元被抽象为PE（处理引擎），其内存层次结构表示为$M_{local}$和$M_{global}$。


图 4. MikPoly概览。

S1：微内核生成（离线）。MikPoly的初始阶段是离线的，它采用模板驱动的调优过程（通过其自动调度组件）来创建和增强微内核。最终生成一组针对特定尺寸优化的微内核。同时，我们为每个微内核开发一个性能模型，使得第二阶段能够以最小的计算开销在线动态选择合适的聚合策略。

S2：微内核聚合（在线）。当一个张量算子的形状在运行时确定后，MikPoly会进入在线的微内核聚合阶段。它的运行时聚合（Runtime Polymerization）组件使用预定义的模式将算子的程序模板重组成不同的实现，并基于一个轻量级的聚合代价模型选择最高效的一个来执行。运行时聚合组件通过将算子模板与预定义模式匹配来派生候选程序，然后使用离线创建的固定大小微内核来实例化它们的参数化微内核，这涉及到启发式地探索运行时形状可用的聚合策略。

3.1 多级加速器抽象

抽象模型定义。MikPoly使用了一个基础的多级加速器抽象来表示现代硬件平台【索引8，Nvidia a100 tensor core gpu: Performance and innovation，2021，IEEE Micro】，【索引35，Ascend: a scalable and unified architecture for ubiquitous deep neural network computing : Industry track paper，2021，2021 IEEE International Symposium on High-Performance Computer Architecture (HPCA)】，【索引36，Cambricon: An instruction set architecture for neural networks，2016，2016 ACM/IEEE 43rd Annual International Symposium on Computer Architecture (ISCA)】，表示为$H = (P_{multi}, M_{local}, M_{global})$。该模型包含多个处理引擎（$P_{multi}$），层次化内存包括单个处理引擎（PE）内的本地内存（$M_{local}$）和多个PE共享的全局内存（$M_{global}$）。这种抽象被当代神经网络编译器如Roller 【索引71，ROLLER: Fast and efficient tensor compilation for deep learning，2022，16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)】、ANT 【索引19，Ant: Exploiting adaptive numerical data type for low-bit deep neural network quantization，2022，2022 55th IEEE/ACM International Symposium on Microarchitecture (MICRO)】和WELDER 【索引52，Welder: Scheduling deep learning memory access via tile-graph，2023，17th USENIX Symposium on Operating Systems Design and Implementation (OSDI 23)】广泛采用，以提高加速器的使用效率。

模型在代价模型中的作用。这个简洁的加速器抽象有效地支持了为性能预测创建精确的代价模型。对于一个给定的张量程序，其在$H$上的并行度依赖于$P_{multi}$，其内存访问特性（独占或共享）由$M_{local}$和$M_{global}$决定。在可能的情况下，使用$M_{local}$来存储数据以提高内存访问效率，而$M_{global}$则将其带宽平均分配给各个PE。在MikPoly中，微内核及其性能模型是针对本地内存$M_{local}$定制的。这种硬件抽象使得MikPoly能够无缝适应不同的加速器，如Nvidia GPU和华为NPU。表1展示了Nvidia A100（$H_{gpu}$）和昇腾910A（$H_{npu}$）的表示。

表 1. $H_{gpu}$ (A100) 和 $H_{npu}$ (昇腾 910A) 的抽象。

3.2 两阶段优化

以GEMM为例说明。我们通过图3中的GEMM示例来详细说明我们为动态形状张量算子创建优化张量程序的方法。

3.2.1 解耦的优化空间

两阶段程序模板。对于一个张量算子（如GEMM），我们通常使用循环分块来增强给定内存层次中的数据复用。我们将其分块程序模板表示为$Q$，它包含一组d维分块循环，带有可调节的分块大小参数。例如，GEMM的程序模板已在图3的❷中展示。

模板结构。与自动调度器【索引7，TVM: An automated end-to-end optimizing compiler for deep learning，2018，13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)】，【索引68，Ansor: Generating High-Performance tensor programs for deep learning，2020，14th USENIX symposium on operating systems design and implementation (OSDI 20)】中使用的传统分块程序模板不同，$Q$体现了一个两阶段结构，包含$Q_{offline}$和$Q_{online}$。其中，$Q_{offline}$是一组为利用$M_{local}$而定制的最内层（离线）循环，而$Q_{online}$是为$M_{global}$优化的其余（在线）循环。这两组循环嵌套分别由图3中❷的蓝色和橙色区域表示。

MikPoly的核心思想。MikPoly的核心思想是从$Q_{offline}$生成不同大小的微内核，并离线优化它们的性能。这使得MikPoly能够根据运行时已知的算子形状，动态地为$Q_{online}$确定最佳的聚合策略。该方法涉及重组$Q_{online}$以创建不同的微内核组合，并由一个精确且轻量级的代价模型指导。

离线优化空间。我们使用一个微内核模板，记作$\hat{K}$，它源自$Q_{offline}$中的离线循环，并为$M_{local}$进行了优化。对于图3所示的GEMM算子，其两阶段模板（❷）产生了一个微内核模板$\hat{K}$（如❸底部所示）。通过使用$\hat{K}$，我们可以利用现有的静态形状自动调度器生成一组优化的固定大小微内核（如❸顶部所示），以及它们的性能模型。这些微内核及其性能模型将在$Q_{online}$的在线聚合过程中使用。

在线优化空间。我们使用预定义的聚合模式重组$Q_{online}$，将$Q$重构为底层算子的不同程序实现。对于GEMM，图3的❹中展示了两种聚合模式。对于每个获得的程序实现，我们通过系统地探索所有可能的聚合策略（本质上是尝试所有离线派生的固定大小微内核）来实例化其参数化微内核，最终选择性能最佳的版本，完成该实现的微内核聚合过程。

3.2.2 优化目标

优化问题定义。给定一个张量算子的两阶段程序模板$Q$和一个在运行时已知的形状，$S_S$表示MikPoly探索的所有张量程序的集合。在硬件平台$H$上为$Q$找到性能最优的程序$S^*$的任务可以定义为一个优化问题：

$$S^* = \arg \min_{S \in S_S} \text{Cost}(S, H)$$

代价模型的作用。由于运行时开销巨大，在运行时于真实硬件上评估$S_S$中的所有张量程序是不切实际的。因此，我们依赖一个聚合代价模型，该模型考虑了并行度、内存访问和资源利用率等因素来估计它们的性能。

3.3 微内核生成

离线阶段。这部分发生在MikPoly的离线阶段。

自动调优固定大小的微内核。MikPoly为每个给定的微内核模板$\hat{K}$生成一个固定大小微内核的集合，记为$S_{\tilde{K}}$。每个微内核$\tilde{K} \in S_{\tilde{K}}$是$\hat{K}$的一个特定尺寸的实例化，经过优化以高效利用给定$H$上的$M_{local}$。图3的❸中展示了一些用于GEMM的固定大小微内核。

生成与筛选过程。MikPoly使用已有的静态形状自动调度器【索引7，TVM: An automated end-to-end optimizing compiler for deep learning，2018，13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)】，【索引68，Ansor: Generating High-Performance tensor programs for deep learning，2020，14th USENIX symposium on operating systems design and implementation (OSDI 20)】为特定平台生成$S_{\tilde{K}}$中的优化微内核。我们使用三个超参数$n_{gen}$、$n_{syn}$和$n_{mik}$，分两步创建$S_{\tilde{K}}$。首先，我们包含所有微内核，每个微内核的嵌套循环来自$\hat{K}$，且每个维度的分块大小来自$\{16 \times i | i \in [1, n_{gen}]\}$。其次，我们只保留一些高性能的微内核，以减少微内核聚合阶段的优化空间。我们使用一个直接从底层算子派生的张量程序，遵循图3中的模式I。我们使用维度大小来自$\{2^j | j \in [0, n_{syn}]\}$生成一组综合测试用例。$S_{\tilde{K}}$中的微内核根据它们在这些综合工作负载上的平均性能进行排名，只保留性能最好的前$n_{mik}$个。

超参数设置。在我们的评估中（第5节），对于所考虑的GPU和NPU平台，我们设置$n_{gen} = 32$，$n_{syn} = 12$，$n_{mik} = 40$。这些经验值覆盖了多样化的真实世界动态形状工作负载，同时最小化了离线自动调优和在线聚合的开销。

微内核性能模型。每个微内核$\tilde{K} \in S_{\tilde{K}}$都有一个由MikPoly创建的性能模型，用于预测其在特定平台$H$上的归约循环中的执行成本。以一个使用尺寸为$(uM, uN, uK)$的微内核$\tilde{K}$的GEMM程序为例，其中$k$是归约循环。GEMM的形状表示为$(M, N, K) = (n_1 \times uM, n_2 \times uN, n_3 \times uK)$。通常，归约循环（$k$）在单个PE上执行，而其余循环（$m$和$n$）则在多个PE上并行化。为了在单个PE上执行归约循环中的$n_3$个$\tilde{K}$实例并重叠计算和内存操作，MikPoly采用了流水线技术【索引44，Pipette: Improving core utilization on irregular applications through intra-core pipeline parallelism，2020，2020 53rd Annual IEEE/ACM International Symposium on Microarchitecture (MICRO)】，【索引71，ROLLER: Fast and efficient tensor compilation for deep learning，2022，16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)】。

流水线任务与并行执行。这种流水线任务可以分为三个阶段：（1）从$M_{global}$加载数据到$M_{local}$；（2）在PE上使用$\tilde{K}$处理$M_{local}$上的数据；（3）将结果从$M_{local}$写回$M_{global}$。执行期间，流水线任务的中间结果存储在$M_{local}$中，减少了内存访问流量。当一个形状为$(M, N, K) = (n_1 \times uM, n_2 \times uN, n_3 \times uK)$的GEMM算子完全执行时，$n_1 \times n_2$个流水线任务（每个任务包含$n_3$个$\tilde{K}$实例）在$P_{multi}$上并行执行。整个算子的执行成本估计为执行$(n_1 \times n_2)/|P_{multi}|$个流水线任务的成本，每个任务由$n_3$个$\tilde{K}$实例组成，其中$|P_{multi}|$表示$H$上$P_{multi}$中的PE数量。

性能模型函数。由于$n_1, n_2, n_3$是在运行时根据具体的GEMM形状确定的，离线阶段只需要为流水线任务创建一个性能模型。我们用$g_{predict}(n, H, \tilde{K})$表示一个分段线性函数，用于估计在平台$H$上执行一个包含$n$个$\tilde{K}$实例的流水线任务的成本。该函数通过在$H$的单个PE上运行$n$从1到$n_{pred}$（经验值设为5120）的$\tilde{K}$实例进行实验，以学习其系数。这些微内核性能模型使得MikPoly能够在其在线微内核聚合阶段高效地估计在$H$的单个PE上执行流水线任务的性能。

3.4 微内核聚合

聚合模式（Polymerization Patterns）。对于给定的程序模板$Q$（例如图3中❷所示的GEMM），MikPoly根据预定义的聚合模式将$Q_{online}$中的在线循环集合划分为多个循环嵌套。这种划分会产生不同的程序实现。每个新形成的循环嵌套都包含来自$Q$的相同微内核模板，但只处理$Q_{online}$原始计算中的特定区域。对于每个这样获得的程序，我们用$R_i$表示其第$i$个循环嵌套（区域）。在GEMM的背景下，图3中可视化了两种这样的模式。为了高效地处理常见场景，我们采用一个模式骨架来系统地生成聚合模式，如图5(a)所示。这个骨架将一个算子的输出划分为七个块，标记为①–⑦。从这个骨架派生出的每个模式都包含多个区域，每个区域包含一个或多个块。为了最小化在线搜索的开销，我们对相似的模式进行分类，并只保留最具代表性的。通过对综合工作负载的评估，我们最终为MikPoly选择了九种独特的代表性模式，如图5(b)所示。例如，图3中展示的模式II将$Q_{online}$分为两个部分：$R_1$ (①–③) 和 $R_2$ (④–⑦)，从而产生两个分别用于micro-kernel(a.uM, a.uN, http://a.uK)和micro-kernel(b.uM, b.uN, http://b.uK)的循环嵌套。


图 5. MikPoly中使用的聚合模式。

聚合策略（Polymerization Strategy）。对于由聚合模式产生的每个程序，MikPoly应用一种聚合策略，从离线生成的固定大小微内核集合中实例化其参数化的微内核。如果一个循环嵌套$R_i$包含一个（参数化的）微内核，其实例化过程就是用$S_{\tilde{K}}$中的一个微内核$\tilde{K}_i$来替换它。此外，MikPoly采用了一种类似于CUTLASS的局部填充技术，以最小化边界检查并保持性能。这确保了对于任何给定形状，都存在带填充的微内核组合。

聚合代价模型。在评估多级加速器$H$上的张量程序$S$的性能时，我们采用以下代价模型。该模型利用为其微内核建立的性能模型，同时考虑了它们并发执行带来的并行性：

$$ \text{Cost}(S,H) = \sum_{(R_i, \tilde{K}_i) \in S} f_{\text{wave}}(R_i, \tilde{K}_i, H) \times f_{\text{pipe}}(R_i, \tilde{K}_i, H) $$

其中，$f_{pipe}$给出了一个微内核流水线执行（一个流水线任务）的成本，$f_{wave}$给出了多个流水线任务并行执行的成本。一个张量程序$S$的总执行成本由执行其各个区域$R_i$的相关成本总和决定，每个区域都包含微内核$\tilde{K}_i$。

并行执行成本函数（f_wave）。函数$f_{wave}$表示并行执行所有流水线任务所需的波次（waves）数量：

$$f_{\text{wave}}(R_i, \tilde{K}_i, H) = \left\lceil \frac{f_{\text{parallel}}(R_i, \tilde{K}_i)}{|P_{\text{multi}}|} \right\rceil$$

其中，$f_{parallel}(R_i, \tilde{K}_i)$表示涉及$R_i$非归约循环的流水线任务（作为$\tilde{K}_i$的实例）的数量。

流水线任务成本函数（f_pipe）。函数$f_{pipe}$用于估计执行一个流水线任务的成本：

$$f_{\text{pipe}} (R_i, \tilde{K_i}, H) = g_{\text{predict}} (f_{\text{num}} (R_i, \tilde{K_i}), \tilde{K_i}, H)$$

其中，$g_{predict}$是在离线阶段获得的性能模型，而$f_{num}(R_i, \tilde{K}_i)$表示在$R_i$的归约循环中出现的$\tilde{K}_i$实例的数量。

3.5 整体流程

算法1概述。算法1概述了MikPoly的工作流程。在离线生成阶段，使用TVM自动调度器【索引7，TVM: An automated end-to-end optimizing compiler for deep learning，2018，13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)】从微内核模板$\hat{K}$生成优化的微内核$S_{\tilde{K}}$（第4行）。在即时聚合阶段，对于运行时已知的动态形状，MikPoly基于两阶段模板$Q$尝试预定义的模式（图5）。利用启发式方法，MikPoly探索聚合策略并使用公式2估计成本（第9-12行）。如果$(R_i, \tilde{K}_i)$的成本超过当前最佳策略的成本，则跳过相关策略，从而在极小的运行时开销下显著缩小搜索空间（第5.3.1节）。最后，MikPoly根据最佳聚合策略构建一个优化的张量程序$S^*$（第13行）。

算法 1 MikPoly的两阶段优化

4 实现

跨平台适应性。尽管GPU和NPU的架构不同，MikPoly的加速器抽象能够统一地表示两者，如表1所示。对于微内核生成，我们根据经验设置超参数来选择要生成的微内核，详见第5.4节。MikPoly采用一个静态形状自动调度器，即TVM，为GPU使用基于CUTLASS的模板，为NPU使用手动模板，以产生固定大小的参数化微内核。这些微内核被编译成二进制文件，保持恒定的形状大小，将张量起始地址和循环迭代次数作为参数供在线确定。

在线聚合过程。在微内核聚合期间，MikPoly为特定的运行时输入形状确定一个合适的聚合策略，并根据可用的运行时数据实例化所选的微内核。这个过程需要调整张量地址偏移量，主要通过标量赋值产生极小的开销。

平台差异化实现。我们为NPU平台采用了九种模式（I - IX），因为在NPU上需要手动指定程序在多个PE（如达芬奇核心）上的并行化。为了将微内核分配给这些核心，我们采用了一种最大最小静态分配算法，以增强并行执行和整体性能。相比之下，在GPU上，由于更强调最小化运行时开销，我们仅限于使用模式I和II。这些模式是基于它们在运行时开销和算子性能之间的最佳平衡来选择的。此外，GPU利用硬件调度器进行动态分配，自动将线程块分配给SM。

开销分析。MikPoly在其离线阶段可以在数小时内（例如，在GPU上为GEMM大约需要6小时）为GPU和NPU平台高效地生成固定大小的微内核。这些针对特定平台定制的微内核，在同一平台上对同一算子无需重新生成。在在线阶段，MikPoly动态选择一个合适的聚合策略，并将偏移量等运行时信息传递给所选的微内核进行分派。主要的运行时开销来自于探索聚合策略和估计它们的成本，这使得MikPoly的运行时开销保持在最低水平。

A4 实验环境

硬件与软件平台。MikPoly的评估涵盖了两个运行Linux操作系统的硬件平台：一个Nvidia A100 GPU和一个昇腾910 NPU（详见表2）。GPU平台使用CUTLASS (v2.9)、CUDA工具包(v11.5)以及cuBLAS和cuDNN库。NPU平台使用CANN SDK (v5.1.RC1)。在GPU平台上，我们使用PyTorch (v1.11)评估CNN模型的端到端性能，使用TurboTransformers（master分支）评估语言模型。在NPU平台上，使用MindSpore (v1.7)评估端到端模型性能。为公平起见，在使用库函数时，我们将卷积转换为GEMM实现。所有实验都经过预热，并取20次运行的平均执行时间。

表 2. 实验平台规格。

基准测试。
- 算子级别：
- GEMM: 涵盖了来自DeepBench【索引43，Deepbench: Benchmarking deep learning operations on different hardware，2016】的166个测试用例和来自真实世界应用（如BERT、DistilBERT、AlexNet等）的1267个测试用例。维度M, N, K中至少有一个是动态的（用*表示）。(见表3)
- Convolution: 涵盖了来自代表性CNN模型（AlexNet, GoogLeNet, ResNet, VGG）的5485个测试用例，涉及不同的滤波器尺寸、特征图尺寸和批处理大小。(见表4)

表 3. 带动态形状的GEMM基准测试。

表 4. 带动态形状的卷积基准测试。

  • 模型级别 (端到端)：
    • 语言模型: HuggingFace的4个模型 (bert-base-uncased, distilbert-base-uncased, roberta-base, albert-xxlarge-v2)。输入为150个长度在5到500之间的句子。
    • CNN模型: TorchVision的4个模型 (alexnet, googlenet, resnet18, vgg11)。使用8种批量大小（$2^i$, i=0..7）和10种分辨率大小（64×j, j=1..10）。
    • 大语言模型 (LLM): HuggingFace的Llama2-13b。在4卡A100 GPU服务器上进行测试。输入序列长度为$2^i$ (i=0..9)，批量大小为$2^j$ (j=0..3)，张量并行度为4。

A4 实验结果

5.2.1 优化动态形状算子

MikPoly vs. GPU库。图6显示，与cuBLAS/cuDNN相比，MikPoly在GEMM上平均加速1.47倍（峰值4.82倍），在卷积上平均加速1.98倍（峰值5.38倍）。与CUTLASS相比，MikPoly在GEMM和卷积上分别平均加速3.02倍和1.72倍。MikPoly在小尺寸形状上表现尤其出色。


图 6. 在GPU上的加速比（以cuBLAS/cuDNN为基准归一化）。

MikPoly vs. NPU库。图7显示，与厂商库CANN相比，MikPoly在NPU上对GEMM平均加速1.10倍，对卷积平均加速1.41倍，在某些测试用例中能通过缓解内存瓶颈实现显著加速。


图 7. 在NPU上的加速比（以CANN为基准归一化）。

5.2.2 优化动态形状模型推理

端到端性能 (GPU)。如图8和图9所示，在语言模型和CNN模型的端到端推理中，MikPoly的性能（包含其运行时开销）显著优于cuBLAS/cuDNN和CUTLASS。
- 语言模型: 对BERT, DistilBERT, RoBERTa, ALBERT分别取得1.39倍, 1.38倍, 1.36倍, 1.37倍的平均加速。
- CNN模型: 对AlexNet, GoogLeNet, ResNet, VGG分别取得1.34倍, 1.69倍, 1.59倍, 1.22倍的平均加速。


图 8. 动态序列长度下的端到端性能比较（以cuBLAS为基准归一化）。


图 9. 动态批量大小和图像分辨率下的端到端性能比较（以cuDNN为基准归一化）。

端到端性能 (NPU)。与CANN相比，MikPoly在AlexNet, GoogLeNet, ResNet, VGG上分别取得1.30倍, 1.19倍, 1.32倍, 1.38倍的平均加速。

5.2.3 与现有技术的比较

MikPoly vs. DietCode/Nimble。在GPU CUDA核上进行公平比较（禁用Tensor Cores）。
- 算子性能: 如图10所示，在1599个GEMM测试用例中，MikPoly平均比DietCode快2.94倍，比Nimble快7.54倍。
- 端到端性能: 如表5所示，在4个语言模型上，MikPoly平均比DietCode快1.55倍。
- 鲁棒性: DietCode要求预定义形状范围，当运行时形状超出范围时会产生错误。如表6所示，在8192个动态形状测试中，MikPoly实现了100%的有效运行，而DietCode的有效运行次数随预定义范围的缩小而大幅减少。即使在形状范围内，MikPoly的性能优势也随着输入范围的扩大而愈发明显（表7）。


图 10. GPU上GEMM的加速比（以DietCode为基准归一化）。

表 5. GPU上的端到端推理性能（以DietCode为基准归一化）。

表 6. 当M的运行时大小超出预定义范围时，比较MikPoly和DietCode的GEMM。

表 7. 当M的运行时大小在预定义范围内时，MikPoly相对于DietCode的GEMM加速比。

5.2.4 在大语言模型(LLM)上的应用

Llama2-13b评估。
- 算子性能: 在Llama2-13b的4个代表性GEMM算子上，相比cuBLAS，MikPoly分别取得了1.09倍、1.24倍、1.21倍和1.08倍的平均加速。（表8）
- 端到端性能: 将MikPoly的GEMM算子集成到Nvidia的FasterTransformer中，与基线相比，对于批量大小为1, 2, 4, 8的情况，分别取得了1.05倍、1.04倍、1.02倍和1.01倍的平均加速。（图11）

表 8. Llama2-13b中GEMM算子的加速比（以cuBLAS为基准归一化）。


图 11. Llama2-13b在GPU上的端到端推理性能（以FasterTransformer为基准归一化）。

5.3 性能分析

在线聚合开销。图12(a)显示，MikPoly的在线聚合成本（橙色部分）仅占总执行时间的一小部分，并且随着形状尺寸的增大，其相对比例会下降，证明了其代价模型的轻量高效。

代价模型有效性。图12(b)比较了MikPoly的不同变体。
- MikPoly-Oracle: 使用穷举搜索找到最优方案，性能最佳但耗时过长（约1.6秒），不实用。
- MikPoly: 综合考虑了并行执行的波次数（$f_{wave}$）和单个流水线任务的执行成本（$f_{pipe}$），其性能非常接近Oracle（达到Oracle的0.96倍），但搜索时间仅为约2微秒。
- MikPoly-Wave (仅考虑波次数) 和 MikPoly-Pipe (仅考虑流水线成本) 的性能则明显较差（分别为Oracle的0.81倍和0.72倍）。
这证明了MikPoly代价模型的有效性和精确性。


图 12. MikPoly在GPU上的GEMM性能分析。

6 案例研究：缓解负载不均衡

问题分析。以GEMM在GPU上(M=4096, N=1024, K=4096)为例，当M维度变化时，使用单一微内核的程序（GEMM-A）会出现性能瓶颈。如图15(a)所示，当M从3328增加到3584时，GEMM-A的执行时间从0.11ms急剧增加到0.21ms。这是因为GPU的线程块是以“波次”（wave）的形式执行的，当所需的线程块总数略大于一个波次能容纳的数量时，最后一个波次的GPU利用率会很低，导致严重的负载不均衡和性能下降（图15(b)）。Nvidia分析工具也证实，此时SM效率从86.67%骤降至58.90%（表9）。


图 14. MikPoly为GPU和NPU上的GEMM使用两种不同聚合方案生成的两个张量程序。

MikPoly的解决方案。MikPoly通过微内核聚合，选择了一个由两个不同微内核（A和B）组成的程序GEMM-AB（图14(a)）。该程序将M=4096的计算任务拆分为M=3072（由A处理）和M=1024（由B处理）两个子任务。通过这种方式，两个子任务的总工作量可以更均匀地填充多个执行波次，避免了最后一个波次利用率过低的问题（图15(c)）。最终，GEMM-AB的SM效率高达96.06%（表9），性能远超GEMM-A。这个案例生动地展示了微内核聚合在解决硬件负载不均衡问题上的有效性。


图 15. MikPoly的微内核聚合如何缓解GPU上的负载不均衡问题（N = 1024, K = 4096）。

表 9. GPU上GEMM的性能测量。

A7 补充细节

7 讨论

通用性。我们已经成功证明MikPoly可以加速GEMM和卷积等代表性张量算子，以及BERT和ResNet等真实应用在GPU和NPU平台上的性能。该框架利用一种新颖的两阶段方法来解决动态形状场景下的性能优化挑战。这个通用框架可以扩展以支持众多其他算子和加速器。

适用性。由于张量形状的多样性，MikPoly在不同应用中实现的加速效果可能有所不同。我们注意到，对于输入形状在较大范围内频繁变化的算子，MikPoly表现得异常出色。此外，如果在编译时已知某些输入形状的范围，我们可以通过生成更合适的微内核集合并优化代价模型来进一步提升性能。

循环变换。本研究中的聚合可以被看作是传统循环变换针对动态形状的一种变体。如图3和图4所示，MikPoly分两个阶段运行，使用一个带有内部离线循环和外部在线循环的张量程序模板。在聚合过程中，MikPoly将在线循环分割成多组嵌套循环，每组都带有一个参数化的微内核，等待循环边界和微内核的选择。选定聚合策略后，这些微内核得以最终确定，同时也设定了嵌套循环的边界。

对LLM系统的影响。MikPoly旨在提升动态形状算子的性能，与in-flight batching技术【索引61，Orca: A distributed serving system for TransformerBased generative models，2022，16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)】完全兼容，支持在运行时动态调整批量大小。这提高了动态形状GEMM算子的效率，从而加速了LLM。未来的计划包括将MikPoly与系统级优化相结合，以进一步提高LLM的效率。

局限性。我们未来的工作将主要集中在两个方向来改进我们的方法。首先，我们计划探索将MikPoly与图级优化技术（如算子融合【索引25，Taso: optimizing deep learning computation with automatic generation of graph substitutions，2019，Proceedings of the 27th ACM Symposium on Operating Systems Principles】）相结合，以在动态形状场景下进一步提升图级的性能。其次，虽然我们目前的实现对卷积采用了基于GEMM的方法，但我们认识到研究其他卷积实现（如Winograd【索引30，Fast algorithms for convolutional neural networks，2016，Proceedings of the IEEE conference on computer vision and pattern recognition】）的潜在好处，这可能带来额外的性能提升。我们期待将这方面作为未来研究工作的一部分进行探索。

8 相关工作

静态形状工作负载优化。对于静态形状的工作负载，研究人员通过自动调优【索引7，TVM: An automated end-to-end optimizing compiler for deep learning，2018，13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)】、多面体模型【索引4，Tiramisu: A polyhedral compiler for expressing fast and portable code，2019，2019 IEEE/ACM International Symposium on Code Generation and Optimization (CGO)】和分析建模等技术在提升算子性能方面取得了成功。相比之下，MikPoly在先前的动态形状自动调度器【索引42，Haotuner: A hardware adaptive operator auto-tuner for dynamic shape tensor compilers，2023，IEEE Transactions on Computers】，【索引51，Nimble: Efficiently compiling dynamic neural networks for model inference，2021，Proceedings of Machine Learning and Systems】，【索引67，Dietcode: Automatic optimization for dynamic tensor programs，2022，Proceedings of Machine Learning and Systems】，【索引72，Disc: A dynamic shape compiler for machine learning workloads，2021，Proceedings of the 1st Workshop on Machine Learning and Systems】中脱颖而出，因为它通过微内核聚合拓宽了优化空间，并采用了两阶段编译方法。这使得MikPoly能够在运行时高效支持任意形状的高性能算子。

图级优化及其他动态编译器。许多研究关注图级优化，包括算子融合【索引25，Taso: optimizing deep learning computation with automatic generation of graph substitutions，2019，Proceedings of the 27th ACM Symposium on Operating Systems Principles】、协同调度和布局选择。在动态形状场景中，DISC【索引72，Disc: A dynamic shape compiler for machine learning workloads，2021，Proceedings of the 1st Workshop on Machine Learning and Systems】使用形状关系而非形状大小作为算子融合的标准。批处理系统采用合并-批处理策略【索引13，Turbotransformers: an efficient gpu serving system for transformer models，2021，Proceedings of the 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming】和请求串联【索引16，Tcb: Accelerating transformer inference services with request concatenation，2022，Proceedings of the 51st International Conference on Parallel Processing】来减少因形状变化造成的填充开销。在IR层面，MLIR【索引29，Mlir: Scaling compiler infrastructure for domain specific computation，2021，2021 IEEE/ACM International Symposium on Code Generation and Optimization (CGO)】和TensorIR【索引15，Tensorir: An abstraction for automatic tensorized program optimization，2022，arXiv】专注于表达性和性能优化。Nimble【索引51，Nimble: Efficiently compiling dynamic neural networks for model inference，2021，Proceedings of Machine Learning and Systems】扩展了Relay以表示动态结构。DietCode【索引67，Dietcode: Automatic optimization for dynamic tensor programs，2022，Proceedings of Machine Learning and Systems】为编译时已知形状范围的动态形状张量算子生成张量程序。VELTAIR【索引39，Veltair: towards high-performance multi-tenant deep learning services via adaptive compilation and scheduling，2022，Proceedings of the 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems】通过多版本编译解决多租户DNN服务器中的资源争用问题。MikPoly则针对DNN系统中不同的动态性问题，专注于优化动态形状的程序，这使得MikPoly可以与其他技术平滑集成，提升整体神经网络性能。

A5 结论

本文介绍了MikPoly，一个新颖的动态形状张量编译器，它为动态形状张量算子优化张量程序。该方法涉及离线创建一组高度优化的固定大小微内核，然后在运行时基于一个轻量级的代价模型，通过微内核聚合在线动态地组合这些微内核。实验结果表明，我们的方法是有效的，在代表性加速器上相比于顶尖的厂商库实现了平均1.49倍的加速。