BOLT: BRIDGING THE GAP BETWEEN AUTO-TUNERS AND HARDWARE-NATIVE PERFORMANCE

文章标题: BOLT：弥合自动调优器与硬件原生性能之间的差距
作者/机构: Jiarong Xing, Leyuan Wang, Shang Zhang, Jack Chen, Ang Chen, Yibo Zhu (机构 1, 2, 3 未在标题页明确区分，但推测来自不同机构)

A1 主要贡献

本文旨在解决当前深度学习（DNN）编译器中自动调优器（如 AutoTVM, Ansor）与设备厂商提供的硬件原生库（如 cuBLAS, cuDNN）之间的性能鸿沟。

核心问题:
1. 性能差距: 自动调优器通过在高层抽象上搜索优化方案，通常难以完全发掘底层硬件的潜力，尤其是在利用如NVIDIA Tensor Cores等专用硬件时，其生成的代码性能远不如厂商手动优化的原生库。例如，Ansor在NVIDIA T4 GPU上执行FP16 GEMM（通用矩阵乘法）的性能仅达到cuBLAS的20%。
2. 调优效率低下: 由于缺乏底层硬件的详细信息，自动调优器的搜索过程效率不高，对一个复杂模型（如ResNet-50）的调优可能需要数小时甚至数天。对于动态模型，依赖缓存调优日志的方法效果不佳。

研究目标:
本文的目标是结合自动调优器的灵活性、可定制性和自动化优势，以及厂商原生库的高性能，实现“两全其美”的效果。

核心创新点 (Bolt):
Bolt利用了设备厂商库（如NVIDIA CUTLASS）日益模块化、模板化和可重构的趋势，提出了一种硬件原生模板化搜索（hardware-native templated search）的新方法，以此为基础带来了三个层面的优化：

  1. 算子层面优化: Bolt设计了一个轻量级的性能分析器，可以直接在硬件原生模板的参数空间中进行高效搜索，从而为张量算子（如GEMM, Conv2D）生成具有硬件原生性能的代码，并将调优时间从数小时缩短至20分钟以内。
  2. 图层面优化: 借助模板库的可扩展性，Bolt实现了新颖的计算图优化。具体而言，它提出了一种名为持久化核函数融合（persistent kernel fusion）的技术，能够将多个连续的GEMM或Conv算子深度融合到单个核函数中，从而减少了中间数据的内存读写和核函数启动开销。
  3. 模型层面优化: Bolt总结了系统与模型协同设计（system-model codesign）的原则，指导研究人员设计对底层优化更友好的神经网络模型。例如，通过选择合适的激活函数、使用1x1卷积加深网络以及对齐张量形状，可以在牺牲少量推理速度的情况下，有效提升模型精度。

通过在TVM和NVIDIA GPU（使用CUTLASS库）上原型实现，Bolt在常用卷积神经网络上的推理速度平均比现有最先进技术（Ansor）提升了2.5倍，并在20分钟内完成了自动调优。

A3 背景知识与动机

自动调优器的性能差距: 当前先进的DNN编译器（如【Chen et al., 2018b; Learning to optimize tensor programs. arXiv, 2018b】和【Zheng et al., 2020a; Ansor: Generating high-performance tensor programs for deep learning. OSDI, 2020a】）利用自动调优来寻找高效的张量实现。它们通过在目标硬件上执行样本程序并观察其性能来构建一个硬件代价模型，然后利用这个模型在巨大的搜索空间中导航。这种远离硬件细节的方法虽然具有平台通用性，但也带来了两个性能问题。

缺乏硬件原生性能: 由于设备模型的不透明性，现有自动调优器生成的张量代码在某些工作负载（例如，非FP32的GEMM/Conv）上与厂商调优库（如cuBLAS和cuDNN）所能达到的硬件原生性能存在差距。如图1所示，Ansor生成的FP16 GEMM代码性能不足cuBLAS的20%，因为它无法有效利用NVIDIA GPU的Tensor Cores这一专用硬件架构。
图1. Ansor（在TVM自动调度器中实现）的速度与cuBLAS可实现的设备速度相比表现明显不佳。工作负载：两个大型方形GEMM和BERT模型中的三个GEMM，批处理大小为32，序列长度为40。

低效的程序搜索: 不透明的设备模型和推断出的硬件执行成本也导致了调优过程的信息不足。现有自动调优器在处理具有多种不同工作负载的模型（如ResNet-152和Inception-V3）时，耗时可达数天或数周【Yu et al., 2021; Lorien: Efficient deep learning workloads delivery. SoCC, 2021】。缓存和重用调优日志（如SAMPL）对静态模型有效，但对具有动态数据结构【Liang et al., 2016; Semantic object parsing with graph lstm. ECCV, 2016】或动态形状【Devlin et al., 2018; Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv, 2018】的模型则效果不佳。相比之下，Bolt通过硬件原生模板化搜索，将常见模型的调优时间缩短到几十分钟。

新兴趋势：模板化库: Bolt的优化之所以成为可能，得益于一个新兴趋势：厂商库正在摆脱早期固定原语的设计，变得模块化和可组合。通过一组声明式参数的控制，模板可以被实例化以适应不同的硬件和工作负载。新的原语可以由现有的原语组合而成，创建新模板的门槛也更低。此外，这些模板化库是考虑了设备细节的高效设计模式，能够提取出不透明自动调优无法达到的硬件性能。

示例：NVIDIA CUTLASS: 我们特别关注的是NVIDIA的模板化库CUTLASS【NVIDIA, b; CUTLASS: CUDA templates for linear algebra subroutines.】。CUTLASS为GEMM的CUDA编程模型的每一层都提供了可重用的C++模板组件。通过设置正确的参数，它能为线程级、warp级、块级和设备级的原语实现高性能。这些模板利用了复杂的设备知识，特别是NVIDIA Volta、Turing和Ampere GPU中集成的Tensor Cores，并为包括B1、INT4、INT8、FP16、BF16、FP32、TF32、FP64、复数和四元数在内的广泛混合精度计算进行了优化。用户通过插入正确的tile大小、数据类型和其他参数，可以为其工作负载挖掘设备级别的性能。除此示例外，Intel的oneDNN【Intel; oneAPI deep neural network library (oneDNN).】和AMD的ROCm【AMD; ROCm libraries.】也展现了类似的设计趋势——带有参数化控制的模板化库。因此，这一设计原则可以推广到其他平台。

CUTLASS的GEMM层次结构: CUTLASS通过将GEMM分解为线程块（threadblock）和warp tile的层次结构，来利用GPU的tiling结构实现高效的GEMM。它优化了数据移动以实现局部性，并仔细控制了从全局内存到共享内存再到寄存器文件的数据移动。图2展示了GEMM操作 C = A · B 中从慢速存储到快速存储的层次结构和数据移动。图2(a)显示了全局内存中的输入A、B和结果C，以及它们的线程块tile。线程块tile可以被划分为共享内存中的warp tile，如图2(b)所示。在这个例子中，一个线程块tile可以被分成八个warp tile，这些warp tile可以进一步被划分为寄存器文件中的线程tile，如图2(c)所示。从全局内存到共享内存再到寄存器文件，内存大小在减小，但读写速度在增加。NVIDIA GPU上的Tensor Cores以线程tile为输入，并将输出存储到寄存器文件中。
图2. CUTLASS中的GEMM层次结构以及在线程块和warp tile中的数据移动。

Bolt：两全其美: Bolt实现了端到端优化，弥合了自动调优器与硬件原生性能之间的差距。

图层面：实现更深度的算子融合: 借助模板化设计，Bolt为算子优化开辟了新的机会。这是因为可以通过模板定制向设备库引入新的优化。Bolt开发了一种名为持久化核函数融合（persistent kernel fusion）的新型算子融合技术，以提高性能。算子融合通过单一核函数计算多个算子，减少数据到内存的 shuffling 以改善局部性【Abdolrashidi et al., 2019; Learning to fuse. NeurIPS ML for Systems Workshop, 2019】，但现有的自动调优器融合方法与高性能的设备库交互不佳。例如，在单个核函数中计算Conv2D+BiasAdd+Hardswish可以提高性能，但像cuDNN这样的固定功能库可能不支持由此产生的算子。通过模板化设计，Bolt启用了考虑更深度融合的新搜索空间，从而开辟了图层面的优化。

算子层面：自动化模板代码生成: 然而，模板化库本身对于普通用户来说过于底层。精确实例化参数以控制tiling大小和数据类型带来了很高的负担。同时，模板原语仅仅是构建块，需要被组装成完整的DNN模型才能执行。Bolt通过将库原语与自动调优器相结合，克服了它们的使用困难。它设计了一个轻量级的性能分析器，以自动搜索最佳的模板参数。通过高效利用硬件细节，该分析器显著缩短了搜索时间。搜索结果随后被用于实例化模板，并生成具有硬件原生性能的底层张量代码。

模型层面：系统友好的模型: Bolt中的端到端优化也为高效的模型设计提供了启示。我们提倡以一种系统友好的方式设计模型，以便它们能高效地利用底层系统提供的优化，从而实现更好的推理性能。在Bolt中，我们总结了三个系统-模型协同设计原则，并通过增强几个RepVGG模型【Ding et al., 2021; Repvgg: Making vgg-style convnets great again. CVPR, 2021】来验证它们。

A2 方法细节

Bolt工作流概述: 图3展示了Bolt的工作流程。它遵循一种“自带代码生成器”（BYOC, Bring Your Own Codegen）【Chen et al., 2021; Bring your own codegen to deep learning compiler. arXiv, 2021】的方法，从张量程序中划分出一个合适的子图，并将其卸载给Bolt进行优化。从用流行框架（如TensorFlow、PyTorch、MXNet）编写的DNN模型开始，Bolt重用TVM前端将模型解析为Relay图。在此图上，它调用计算图优化（如更深度的融合），并在优化后的图上执行图分区。接下来，Bolt执行硬件原生性能分析，为Bolt子图中的每个算子搜索最佳的核函数。最后，Bolt生成高性能的CUDA代码，这些代码将与TVM生成的代码一起编译成一个单一的运行时文件。
图3. Bolt的工作流程。蓝色方框是我们的贡献。

3.1 实现更深度的算子融合

通过扩展硬件原生模板进行图级优化: Bolt通过扩展硬件原生模板，实现了新颖的图级优化。具体来说，Bolt引入了“持久化核函数”（persistent kernels），这使得新颖的、更深度的算子融合成为可能。如图4(a)所示，它以epilogue融合为基础，并进一步融合两个或多个连续的GEMM/Conv。将多个GEMM或Conv融合成一个单一的算子可以通过以下方式提高性能：(i) 消除存储和加载层间激活值的内存流量；(ii) 消除启动延迟，这对于小批量大小的短核函数尤其有益；(iii) 扩大编译器的优化范围，以探索更多的指令调度选项【Wu et al., 2012; Kernel weaver: Automatically fusing database primitives for efficient gpu computation. MICRO, 2012】。图4(b)展示了持久化核函数融合的核函数视图。
未融合:
未融合的图视图
Epilogue融合:
Epilogue融合的图视图
持久化核函数融合:
持久化核函数融合的图视图
(a) 持久化核函数融合的图视图

前提条件：Epilogue融合: 作为持久化核函数融合的前提，Bolt首先集成了CUTLASS中提供的epilogue融合功能，该功能将一个GEMM/Conv核函数与其后的epilogue（收尾操作）融合在一起，形成一个单一算子，以便我们能进一步利用持久化核函数。CUTLASS中的epilogue融合模式包括：(i) 逐元素算子，(ii) 数据类型转换，(iii) 列上广播向量，以及(iv) 列上部分归约。Bolt在计算图中识别这些模式，并自动生成相应的算法策略。Bolt以epilogue融合为起点，利用持久化核函数开发了更深度的融合。
图4. 连续GEMM/Conv的持久化核函数融合的图视图和核函数视图。

3.1.1 持久化核函数（GEMM/Conv）融合

持久化核函数融合的原理: 持久化核函数允许将多个GEMM或Conv融合成一个算子以提高性能。如图4(b)所示，当两个GEMM/Conv操作被融合在一起时，这两个算子的数学计算主循环在一个融合后的单一核函数中背靠背运行。第一个GEMM/Conv的输出激活值保留在更快的GPU内存中。这消除了将GEMM0/Conv0的输出激活值存回全局内存、启动GEMM1/Conv1核函数、以及从全局内存加载GEMM1/Conv1输入激活值的需要。Bolt自动识别使用持久化核函数的机会，并通过在CUTLASS中创建新模板来生成CUDA代码。下面将详细描述背靠背GEMM融合，卷积融合的工作方式类似。

背靠背GEMM的定义与融合条件: 一个背靠背GEMM定义如下：
$D0 = \alpha_0 A0 \cdot W0 + \beta_0 C0,$
$D1 = \alpha_1 D0 \cdot W1 + \beta_1 C1,$
其中A0、W0和W1是矩阵输入，αs和βs是标量输入，C0、C1是预先存在的矩阵（偏置），它们将被输出覆盖。为了融合背靠背的GEMM，第一个GEMM层的输出激活值D0必须用作第二个GEMM层的输入激活值。这要求所有层的GEMM的M维度保持不变。对于背靠背的Conv，这要求所有后续的Conv（从第二个开始）必须使用1×1的滤波器，没有填充，并且步长为1。

关键属性：线程块驻留（Threadblock residence）: 持久化核函数的关键挑战是在不从全局内存加载其输入激活值的情况下计算第二个GEMM/Conv。这要求第一个GEMM/Conv的每个输出线程块都保留在同一个线程块的内存中（无论是在共享内存还是寄存器文件中），作为其各自的输入线程块。我们称之为线程块驻留。如果这个条件不成立，第二个GEMM/Conv就必须从全局内存中获取数据，从而消除了持久化核函数的优势。对于GEMM融合，线程块驻留要求每个算子的ThreadBlock N = GEMM N。至于Conv融合，要求是ThreadBlock N = Conv output channel。图5可视化了这一要求。基于线程块驻留，我们为不同场景开发了两种设计。
图5. GEMM融合的线程块驻留示意图。彩色框代表一个单一的线程块。这要求ThreadBlock0 N = N0，ThreadBlock1 N = N1。

RF驻留融合（RF-resident fusion）: 当权重矩阵W1可以在其'N'维度上完全流式传输到一个warp tile中时（如图6所示），可以通过将每个线程块的输出激活值完全存储在寄存器文件（RF）中来满足线程块驻留。这样做，第二个GEMM/Conv可以在不接触其他warp的情况下计算W1。我们称之为RF驻留融合，它要求每个层的warp大小必须遵循Warp N = ThreadBlock N = GEMM N。在RF驻留融合中，每个warp将在RF中拥有一块由当前层产生的累加器数据（称为累加器片段）。这将完全用作由同一warp计算的下一层的输入。我们开发了一个CUTLASS warp fragment iterator来从累加器片段中提取数据，并将其馈送到warp级的MMA操作中。RF驻留融合通过扩展CUTLASS中线程块级的GEMM设计，整合了背靠背的MMA流水线。我们的设计在GEMM操作之间没有干扰。第二个GEMM唯一的额外操作是从前一个累加器获取warp片段，并在RF中执行epilogue计算。
图6. 背靠背GEMM线程块中的RF驻留融合。线程块和warp大小的要求是：Warp0 N=ThreadBlock0 N=N0, Warp1 N=ThreadBlock1 N=N1。

共享内存驻留融合（Shared memory-resident fusion）: RF驻留的GEMM融合会产生较高的RF压力，尤其是在GEMM N较大时，这可能会损害核函数性能并限制其适用场景。为了解决这个问题，我们提出了共享内存驻留融合，以放宽warp大小的限制。在这种设计中，当第二个GEMM/Conv需要warp之间共享数据时，数据可以被暂存到共享内存而不是RF中。图7展示了一个例子，其中D1的计算必须从'N'维度上的多个warp tile中流式传输W1片段。因此，在GEMM0中产生的累加器数据必须从RF传输到共享内存，以便被GEMM1加载。每个warp拥有的数据块将在M维度上为下一层共享。通过这样做，可以放宽RF驻留融合中Warp N的大小限制。为了实现共享内存驻留融合，我们引入了一个smem fragment iterator作为将累加器tile存储到共享内存，然后从共享内存为第二个GEMM获取片段的机制。为了获得更高的性能，我们精心设计了共享内存布局，以避免在存储第一个核函数的累加器和为第二个核函数加载它时出现任何共享内存库冲突。
图7. 背靠背GEMM线程块中的共享内存驻留融合。线程块大小的要求是：ThreadBlock0 N = N0 ≠ Warp0 N, ThreadBlock1 N = N1 ≠ Warp1 N。

融合技术总结: RF驻留和共享内存驻留融合实现了连续GEMM/Conv的更深度融合。基于背靠背融合，Bolt可以通过扩展持久化核函数模板和复制GEMM流水线来支持融合多个GEMM/Conv。

3.2 自动化模板代码生成

3.2.1 代码生成中的挑战

模板化库带来的新挑战: 模板化库为端到端张量程序优化带来了新的挑战。首先，这些模板通常不为端到端模型提供完整的功能，只支持一部分算子。一个简单的解决方案是为每种硬件从头开发一个完整的编译器栈，但这不具备可扩展性。Bolt通过采用BYOC（Bring Your Own Codegen）【Chen et al., 2021; Bring your own codegen to deep learning compiler. arXiv, 2021】方法解决了这一挑战。这使我们能够尽可能地重用现有的编译器栈（如TVM），并专注于使用模板化设备库进行优化和代码生成。

朴素BYOC方法的局限性与Bolt的解决方案: 单纯应用BYOC并不能解决所有问题。首先，模板化设备库本身不能直接运行，需要用户使用经过良好调优的参数来实例化模板以获得良好性能，但BYOC不支持这种性能分析。Bolt通过提出一个轻量级的硬件原生性能分析器解决了这个问题，该分析器可以在几分钟内为特定工作负载的算子搜索到最佳参数。此外，传统的BYOC将设备库视为不可知的外部函数，并生成带有钩子的硬件代码以便在运行时调用它们。这种设计使得定制硬件库和支持新的优化（如布局转换和核函数填充）变得困难。Bolt通过将库视为白盒并直接以其约定生成代码来解决这个问题。

3.2.2 轻量级性能分析器

高效搜索最佳模板参数: Bolt设计了一个轻量级的性能分析器来搜索最佳的模板参数。传统的自动调优器假定没有硬件信息，通过生成样本实现并测量其速度来推断成本模型，这需要自动调优器探索一个巨大的搜索空间，导致调优时间很长。Bolt通过将耗时的样本程序生成与性能测量分离开来，并有效利用硬件细节来加速前者，从而大大减少了搜索时间。具体来说，CUTLASS模板中与性能相关的参数包括线程块、warp和指令形状、swizzling functor、stages等。得益于白盒方法，Bolt根据GPU架构以及每个硬件特有的调优指南来确定它们的可能值。例如，在寄存器文件容量内，Bolt偏好使用大的warp tile尺寸以获得更高的计算-内存比；在现代NVIDIA GPU上运行时，每个线程块使用四个或八个warp往往性能更好；小问题尺寸需要小的线程块尺寸来启动足够的线程块以保持更多SMs繁忙。对于每个GPU架构，Bolt会产生数十个最佳参数组合，并通过初始化模板生成相应的样本程序。请注意，这些样本程序通过给定不同的输入，可以在模型和工作负载之间重用。因此，在运行时，Bolt可以通过使用具体输入调用预先生成的样本程序来分析性能。

3.2.3 模板化代码生成

白盒式代码生成: 传统的BYOC系统【Chen et al., 2021; Bring your own codegen to deep learning compiler. arXiv, 2021】无法以模板化格式为目标进行代码生成；它们在运行时将此类库视为外部函数。相反，Bolt通过使用分析器识别出的最佳参数实例化模板，直接以CUTLASS的约定生成底层张量实现。我们的方法有两个优点。首先，生成的代码提供了卓越的性能，例如，在Ampere A100上FP16 GEMM的吞吐量可以达到300 TFLOPS，超过了硬件理论极限的95%。其次，它为在生成的代码中添加新颖的优化提供了充分的灵活性。在Bolt中，我们开发了以下两种优化。

布局转换: CUTLASS只支持Conv的NHWC布局，因为它比NCHW布局更快【NVIDIA, c; Tensor layouts in memory: NCHW vs NHWC.】。但并非所有模型都以期望的布局编写——例如，所有的PyTorch模型都使用NCHW。为了创造更多的优化机会，Bolt提供了自动布局转换。请注意，这与TVM中通过修改Relay图来实现功能的布局转换不同。相反，Bolt直接在模型第一层和最后一层的生成CUDA代码中实现转换，以节省额外的核函数启动开销。该转换需要一个新的张量来保存新布局的数据。在核函数内部分配和初始化新张量会产生显著的内存开销。相反，我们通过在模型的参数中添加一个新变量来预先分配内存，这个变量可以直接被核函数使用。

核函数填充: 尽管CUTLASS支持对齐方式8、4、2、1以覆盖所有不同的工作负载，但不同对齐方式的性能差异很大。NVIDIA GPU支持的最大向量化加载和存储是128位，因此对于FP16数据类型，最有效的使用方式是对齐8（128/16）。在这种情况下使用对齐8可以减少加载和存储指令的数量，以及每次加载和存储指令所需的谓词数量。张量形状的维度如果不能被8整除，就必须使用较小的对齐方式。例如，卷积神经网络的第一层通常有三个输入通道，这必须使用对齐1。这将导致非合并的内存访问和共享内存库冲突。因此，Bolt会自动填充未对齐的张量以使用对齐8。这使我们不仅能充分利用tensor core加速，还能减少内存加载时间。与布局转换类似，我们也在模型的参数中预先分配对齐后的张量内存。

3.3 设计系统友好的模型

从系统优化到模型设计: Bolt中的图级优化（如持久化核函数融合）和算子级优化（如自动填充）也为模型级的优化机会提供了启示。以一种能有效利用系统优势的方式设计的模型可以带来更高效的推理。我们称之为系统-模型协同设计，它可以帮助构建运行更高效的系统友好模型。Bolt确定了以下协同设计原则。

利用Epilogue融合探索不同的激活函数: 激活函数的选择对DNN模型的准确性有显著影响【Prajit et al., 2017; Swish: a self-gated activation function. arXiv, 2017】。多年来，设计了一系列激活函数，如ReLU【Nair & Hinton, 2010; Rectified linear units improve restricted boltzmann machines. ICML, 2010】及其变体、GELU【Hendrycks & Gimpel, 2016; Gaussian error linear units (gelus). arXiv, 2016】、Softplus【Zheng et al., 2015; Improving deep neural networks using softplus units. IJCNN, 2015】和Hardswish【Howard et al., 2019; Searching for mobilenetv3. ICCV, 2019】。在Bolt中，epilogue融合会将激活函数与前面的GEMM/Conv融合，以减少激活函数的开销。因此，模型设计者可以在他们的模型中探索不同的激活函数，并确定最有效的一个。

使用1×1 Conv加深模型: 加深神经网络以获得更高的准确性是一种常用的模型设计技术。例如，ResNet【He et al., 2016; Deep residual learning for image recognition. CVPR, 2016】有从18层到151层的不同深度，准确性也随之增加。然而，随着深度的增加，推理速度会迅速下降。另一方面，在Bolt中使用1×1 Conv加深模型只会产生较低的计算开销。这是因为有持久化核函数融合优化。因此，尽管使用1×1 Conv加深模型在提高准确性方面不如使用更大的核，但人们仍然可以添加1×1 Conv来提高准确性，同时减少速度损失。

对齐张量形状以更高效地使用GPU: 正如我们在3.2.3节中讨论的，张量形状对模型在GPU上运行的效率有重大影响。尽管Bolt会自动对未对齐的张量进行填充，但填充本身会产生额外的开销，如表3所示。因此，人们可以设计具有对齐张量形状的模型以实现更高的效率，避免需要额外的填充。

A4 实验环境

  • 硬件配置:
    • GPU: 单张 NVIDIA Tesla T4
  • 软件配置:
    • 基线系统: Ansor，TVM 中最先进的自动调优器。
    • 数据类型: 所有推理计算均使用 FP16。
  • 模型与数据集:
    • 微基准测试:
      • GEMM: 来自 BERT 模型的典型 GEMM (batch size=32, seq len=40) 和两个方形 GEMM。
      • Conv2D: 来自 ResNet-50 的 Conv2D (batch size=32)。
      • 融合测试: 来自真实推荐模型（如 DCNv2, DLRM）的 GEMM 和 RepVGG 模型的 Conv2D。
    • 端到端测试:
      • 模型: VGG, ResNet, RepVGG 等六个广泛使用的卷积神经网络。
    • 案例研究:
      • 模型: RepVGG-A0, A1, B0。
      • 数据集: 在 ImageNet 上进行训练。

A4 实验结果

4.1 微基准测试

  • GEMM/Conv2D 性能 (图8):

    • 实验内容: 对比 Bolt 生成的 GEMM 和 Conv2D 实现与 Ansor 的性能。
    • 实验结果: 在计算密集型 GEMM 上，Bolt 比 Ansor 快 6.1-9.5 倍；在计算不那么密集的 GEMM 上快 1.9 倍。在所有 Conv2D 测试用例中，Bolt 比 Ansor 快 2.7-3.5 倍。
    • 结论: Bolt 基于硬件原生模板的调优策略能有效发掘硬件原生性能，性能显著优于 Ansor。
      图8. Bolt在GEMM和Conv2D上的性能。图8a显示了BERT中批大小为32、序列长度为40的GEMM和两个方形GEMM的速度。图8b显示了ResNet-50中3 × 3 Conv2D的速度。批大小为32，所有Conv2D都使用(1, 1)零填充。

  • Epilogue 融合性能 (图9):

    • 实验内容: 评估 GEMM/Conv2D+BiasAdd+Activation 模式的 Epilogue 融合效果，测试了 ReLU, GELU, Hardswish, Softplus 四种激活函数。
    • 实验结果: 与 Bolt 不进行 Epilogue 融合（即由 TVM 融合 BiasAdd 和激活）相比，Epilogue 融合使 GEMM 的平均速度提升了 1.45 倍，Conv2D 的平均速度提升了 1.38 倍。
    • 结论: Epilogue 融合能有效提升包含逐元素操作的计算模式的性能。
      图9. Epilogue融合在GEMM/Conv2D+BiasAdd+Activation模式上的性能。GEMM的工作负载为M=1280, N=3072, K=768。Conv2d的工作负载为H=W=56, IC=OC=64, kernel=(3, 3), stride=(1,1), padding=(1,1)。

  • 持久化核函数融合性能 (表1, 表2):

    • 实验内容: 评估融合两个背靠背 GEMM 和两个背靠背 Conv2D 的性能。
    • 实验结果: 对于 GEMM，持久化核函数融合带来了 1.2x-1.5x 的加速（表1）。对于 Conv2D（一个 3x3 Conv2D 跟一个 1x1 Conv2D），性能提升了 1.1x-2.0x（表2）。
    • 结论: 持久化核函数融合能通过减少内存访问和核函数启动来显著提升连续 GEMM/Conv 的计算速度。
      表1. 使用持久化核函数融合两个背靠背GEMM的性能。每个GEMM后跟一个ReLU epilogue，所有这些都将融合成一个核函数。
      表1. 使用持久化核函数融合两个背靠背GEMM的性能。
      表2. 使用持久化核函数融合两个背靠背Conv2D的性能。每个Conv2D后跟一个BiasAdd和一个ReLU epilogue。3 × 3 Conv2D使用(1, 1)填充，1 × 1 Conv2D使用(1, 1)步长且无填充。
      表2. 使用持久化核函数融合两个背靠背Conv2D的性能。

  • 填充性能与开销 (表3):

    • 实验内容: 评估 Bolt 的自动填充功能对输入通道数不能被 8 整除的 Conv2D 的性能提升及其开销。
    • 实验结果: 通过填充使对齐方式从 2 变为 8，计算速度平均提升了 1.8 倍。然而，填充操作本身占总计算时间（填充+卷积）的平均 16%。
    • 结论: 自动填充能带来显著性能收益，但也引入了不可忽视的开销，这印证了“设计具有对齐张量形状的模型”这一协同设计原则的重要性。
      表3. Bolt自动填充的性能和开销。未填充的Conv2D使用alignment=2计算；填充后，可使用alignement=8。填充成本是填充时间占总计算时间（填充+Conv2D）的比例。
      表3. Bolt自动填充的性能和开销。

4.2 端到端优化

  • 实验内容: 在六个广泛使用的卷积神经网络上评估 Bolt 的端到端性能，并与 Ansor 对比。
  • 实验结果 (图10):
    • 推理速度: Bolt 的推理性能显著优于 Ansor。在 VGG 模型上快 4.2 倍，ResNet 模型上快 1.5 倍，RepVGG 模型上快 2.6 倍，平均速度提升了 2.8 倍。
    • 调优时间: Bolt 对所有模型的调优都能在 20 分钟内完成，而 Ansor 平均需要 12 小时。
  • 结论: Bolt 在端到端模型优化上实现了性能和效率的双重飞跃，证明了其方法的优越性。
    图10. 常用卷积神经网络的归一化推理速度和调优时间。

4.3 系统友好模型：RepVGG 案例研究

  • 实验内容: 应用系统-模型协同设计原则来增强 RepVGG 模型。
  • 实验结果:
    • 更换激活函数 (表4): 在 RepVGG-A0 中，使用 Hardswish 替代 ReLU，Top-1 准确率提高了 0.67%，而推理速度仅下降了 3.3%，证明了探索不同激活函数的价值。
    • 使用 1x1 Conv 加深模型 (表5): 在每个 3x3 Conv 后增加一个 1x1 Conv（Bolt会进行融合），模型准确率平均提升了约 0.8%，而速度平均仅下降 15.3%，是一种高效提升精度的手段。
    • 综合效果 (表6): 结合以上两种技术（1x1 Conv + Hardswish）并使用高级训练策略，增强后的 RepVGG 模型在精度提升上比原始模型的扩展方式更高效。例如，RepVGGAug-A1 相较于 RepVGG-A1，在相似的速度开销下，精度提升了 1.83%，远高于 RepVGG 从 A1 扩展到 B0 带来的 1% 精度提升。
  • 结论: 系统-模型协同设计原则是有效的，能够指导研究人员在模型准确率和推理速度之间做出更优的权衡，设计出更高效的模型。
    表4. RepVGG-A0使用不同激活函数的top-1准确率和速度（120个epoch + 简单数据增强）。
    表4. RepVGG-A0使用不同激活函数的top-1准确率和速度。
    表5. 原始RepVGG模型及其使用1 × 1 Conv2D增强后的top-1准确率和速度（200个epoch + 简单数据增强）。
    表5. 原始RepVGG模型及其使用1 × 1 Conv2D增强后的top-1准确率和速度。
    表6. 原始RepVGG模型及其使用1 × 1 Conv2D+Hardswish增强后的top-1准确率和速度（300个epoch + 高级增强、标签平滑和mixup）。
    表6. 原始RepVGG模型及其使用1 × 1 Conv2D+Hardswish增强后的top-1准确率和速度。

A7 补充细节

5 讨论

其他平台: 尽管本文使用NVIDIA CUTLASS和GPU来展示Bolt的设计，但我们的方法并不局限于任何特定的设备或库。模板化库是其他设备的设计趋势。我们期望对于其他库，类似Bolt的方法也能带来新的端到端优化。

持久化核函数融合的局限性: 尽管持久化核函数可以融合任何遵循线程块驻留原则的顺序GEMM/Conv2D，但我们专门为内存受限的算子设计它，这与通用算子融合的动机一致。也就是说，Bolt可以提高具有小N和K维度但大M维度的顺序GEMM以及具有小通道数的Conv2D的性能。融合计算受限的算子可能会因为线程块驻留的要求而导致性能下降。

6 相关工作

自动调优器: 许多DNN编译器和框架采用自动调优策略来搜索性能最优的张量实现（如【Chen et al., 2018b; Learning to optimize tensor programs. arXiv, 2018b】, 【Zheng et al., 2020a; Ansor: Generating high-performance tensor programs for deep learning. OSDI, 2020a】, 【Adams et al., 2019; Learning to optimize halide with tree search and random programs. TOG, 2019】, 【Zheng et al., 2020b; Flextensor: An automatic schedule exploration and optimization framework for tensor computation on heterogeneous system. ASPLOS, 2020b】）。由于它们通过尝试不同的张量实现并测量其性能来推断硬件成本模型，这需要数小时到数天的时间。此外，生成的张量程序无法达到硬件原生性能。Bolt弥合了自动调优器和硬件原生性能之间的差距。

算子融合: 算子融合是一项重要的图级优化（如【Chen et al., 2018a; TVM: An automated end-to-end optimizing compiler for deep learning. OSDI, 2018a】, 【Abadi et al., 2016; Tensorflow: A system for large-scale machine learning. OSDI, 2016】, 【Leary & Wang, 2018; XLA - tensorflow, compiled, 2018】, 【Paszke et al., 2019; Pytorch: An imperative style, high-performance deep learning library. NeurIPS, 2019】, 【Abdolrashidi et al., 2019; Learning to fuse. NeurIPS ML for Systems Workshop, 2019】, 【Jia et al., 2019; Taso: optimizing deep learning computation with automatic generation of graph substitutions. SOSP, 2019】）。然而，现有的算子融合只考虑一个GEMM/Conv及其相邻的算子，例如BiasAdd、ReLU，并且这种融合并未得到厂商库的良好支持。Bolt实现了具有高性能的新型算子融合。例如，提出的持久化核函数融合可以融合一系列GEMM和Conv，进一步提高性能。我们的持久化核函数不同于Persistent RNNs【Diamos et al., 2016; Persistent rnns: Stashing recurrent weights on-chip. ICML, 2016】，后者是专门为RNN手动设计的，且不使用tensor core。

系统友好的模型设计: RepVGG【Ding et al., 2021; Repvgg: Making vgg-style convnets great again. CVPR, 2021】通过在训练模型时采用多分支架构追求高准确率，并通过重参数化移除分支以进行推理，从而设计出系统友好的模型。此外，RepVGG仅使用3×3的Conv2D，这种算子得到了硬件的良好支持。Bolt通过提出系统-模型协同设计原则进一步扩展了这一思想，并以RepVGG作为具体案例研究。

A5 结论

本文提出了Bolt，一个旨在弥合自动调优器与设备库原生性能之间差距的系统。Bolt利用了厂商库正变得模块化和可组合的新兴趋势，它结合了自动调优器的灵活性和模板化设备库的硬件原生性能，从而实现了两者的优点。我们的设计带来了新的张量级和图级优化，并启发了系统友好的模型设计见解。实验表明，与现有最先进技术相比，Bolt可以在广泛使用的卷积神经网络上实现2.5倍的加速。此外，它在20分钟内即可完成自动调优。