BOLT：弥合自动调优和硬件原生性能之间的差距

本文介绍的BOLT基于TVM框架，在GPU平台上进行了进一步的图优化和算子优化，最终将常见的卷积神经网络模型的推理速度提升了2.5倍，搜索时间大大缩减，可以实现20分钟内完成自动搜索调优。下面对BOLT的主要技术细节进行介绍。

1 背景

NVIDIA CUTLASS 是一个开源项目，是 CUDA  C++ 模板抽象的集合，用于在 CUDA 内部实现高性能矩阵乘法和卷积运算。它定义了一系列高度优化的算子组件，开发人员可以通过组合这些组件，开发出性能和 CUDNN、CUBLAS 相当的线性代数算子。CUTLASS 通过将计算拆分为 thread block tiles、warp tiles 和 thread tiles，高效地实现了GPU 中的矩阵乘法。如图 1 所示，可以看到数据从全局内存移动到共享内存，从共享内存移动到寄存器，从寄存器移动到 SM CUDA Cores 进行计算。

图1 CUTLASS 中 GEMM 计算的拆分过程

目前，在GPU平台上，TVM使用自动搜索得到的性能与 CUTLASS 算子库的性能存在很大的差距。TVM 需要一个 tophub 数据库来存储各种shape的调优调度，当给定一个以前没有搜索过的shape时，将会花费几个小时去搜索或者转到性能很差的默认调度；此外，针对不同数据类型，TVM auto scheduler 不能良好支持NVIDIA Tensor Core 指令。这些问题阻碍了 TVM 在真实的推理部署场景中落地应用，针对以上情况，BOLT提出将 CUTLASS 引入到 TVM Codegen 中，并利用其算子融合能力，大幅提升模型性能。据论文作者所述，BOLT 已经部署在真实生产环境中，并且代码已经合入到TVM项目中。

2 BOLT设计

图 2 显示了 BOLT 的整体架构。BOLT基于TVM框架，采用BYOC（Bring Your Own Codegen）方法，将CUTLASS引入到TVM。主要创新点有两个，一是图级优化——Deeper fusion，深度学习模型导入TVM，转化为relay IR后，进行更深层次的算子融合；二是算子级优化——Perf profiler，单纯的使用BYOC引入CUTLASS并不会直接得到最好结果，通过设计一个性能分析器来搜索最优参数，从而得到最优性能。

图2 Bolt 的工作流程（蓝框是创新点）

3 架构详解

3.1 Deeper fusion

BOLT 提出了 Persistent kernel 的概念，实现了更深层次的算子融合。如图 3（a）所示，以 Epilogue fusion 为基础，进一步融合了两个连续的 GEMMs/Convs。图 3（b）显示了 Persistent kernel 融合的kernel视图，通过消除加载、存储中间层激活函数的数据移动时间和启动时延来提高性能。两个算子计算中的主循环在融合后的内核中相继执行。BOLT 自动识别使用 Persistent kernel 的时机，GEMM融合计算定义为：

A0、W0 和 W1 为矩阵输入，和为标量输入，C0 和 C1 为 偏置。为了实现更深层次的融合，第一个 GEMM 的输出 D0 必须用作第二个 GEMM 的输入。这要求 GEMM 的 M 维对于所有层都保持相同， 对于Convs，要求后续的 Convs 必须使用 1 × 1 卷积核，没有填充且步幅为 1。

图3 GEMM/Convs Persistent kernel 融合的 graph 视图和 kernel 视图

Persistent kernel的关键挑战在于不从全局内存加载输入激活的情况下计算第二个 GEMM/Conv。要求第一个 GEMM/Conv 的每个输出 threadblock 与第二个 GEMM/Conv 的输入 threadblock 保存在相同的 threadblock 内存中（在共享内存或寄存器中）。如图 4 所示，对于 GEMM 融合，需满足每个算子的 ThreadBlock_N = GEMM_N。对于 Conv 融合，需满足 ThreadBlock_N = Conv 输出通道。针对不同的场景有两种设计。

图4 GEMM 融合的 threadblock 驻留示意图

RF-resident 融合。当权重矩阵 W1 的 N 等于 warp tile 时（如图 5 所示），可以将每个 threadblock 的输出激活完全存储在寄存器中 (RF) 。这样使得第二个 GEMM/Conv 可以在不影响 W1 的其他 warps 的情况下进行计算，这种方式称为 RF-resident 融合，它要求每层的 warp 大小必须遵循 Warp_N = ThreadBlock_N = GEMM_N。在 RF-resident 融合中，每个 warp 将拥有当前层产生的 RF 中的一个累加器数据块（称为累加器片段）。这将用作由同一 warp 计算的下一层的输入。BOLT开发了一个 CUTLASS warp 片段迭代器来从累加器片段中提取数据并进行 WMMA 运算。这种设计对 GEMM 运算没有干扰。第二个 GEMM 唯一的额外操作是从前一个累加器中获取 warp 片段，并在寄存器中执行最终计算。

图5 RF-resident 融合

Shared memory-resident 融合。当 GEMM 的 N 很大时，GEMM的 RF-resident 融合会产生更高的寄存器压力，这可能会损害内核性能并限制使用场景。为了解决这个问题，BOLT提出了共享内存驻留融合来放宽warp大小的限制。在这个设计中，当第 2 个 GEMM/Conv 需要在 warp 之间共享数据时，可以将数据暂存到共享内存。如图 6 所示，D1 的计算必须从多个 wrap 块中流式传输 W1 片段。因此，GEMM0 中产生的累加器数据必须从寄存器传输到共享内存，才能被 GEMM1 加载。每个 warp 拥有的数据块在 M 维中共享给下一层。为了实现共享内存驻留融合，BOLT引入了一个 SMEM 片段迭代器作为将累加器块存储到共享内存中的机制，然后从共享内存中获取片段以用于第二个 GEMM。为了获得更高的性能，BOLT还设计了共享内存布局，以避免发生任何共享内存库冲突。

图6 Shared memory-resident 融合

3.2 Perf profiler

由于CUTLASS模板库是针对于单个算子实现的，并不能直接在模型中应用，BOLT通过使用TVM BYOC的方法，将CUTLASS引入到Codegen。但仅使用BYOC，并不能获得最优性能，需要用户找到最优参数实例化模板才可实现，过程需要大量的专家经验和时间。BOLT 通过提出一种轻量级的硬件原生性能分析器解决了这个问题，该分析器可以在几分钟内为具有特定工作负载的算子搜索出最佳参数。

传统的自动调优器，通过生成样本并测量其速度来推断 Cost model，这需要大量的搜索空间和较长的调优时间。BOLT 通过将耗时的样例程序生成与性能测量分离，并通过有效利用硬件细节进行加速，大大减少了搜索时间。CUTLASS 模板中与性能相关的参数包括 threadblock、warp 和 instruction shapes、swizzling functor和stage 等。BOLT 采用白盒方法， 根据 GPU 架构以及特定于每个硬件的调优指南确定它们的可能值。例如，在寄存器的容量范围内，BOLT 更喜欢使用较大的wrap tile 尺寸，以实现更高的计算内存比率；在现代 NVIDIA GPU 上，每个线程块 4 个或 8 个 warp 往往具有更好的性能；小问题需要小线程块来启动足够的 thredablocks 来保持更多 SM 忙碌。对于每一个GPU架构，BOLT 都会产生数十个最佳参数组合，并通过初始化模板生成对应的示例程序。这些样例程序可通过给定的不同输入跨模型和工作负载重用。因此，在运行时，BOLT 可以通过调用带有具体输入的预生成示例程序来分析性能。

4 评估

4.1 微基准测试

GEMM 选择 BERT 中的典型 GEMM（其中 batch size 为 32，sequence length 为 40 ）和两个方形 GEMM ，作为工作负载。以 TVM Ansor 作为基准，按照 TVM 官方示例调整每个工作负载进行 2000 次试验。运行每个工作负载 1000 次并计算平均速度；结果如图 7（a）所示。BOLT 在计算密集型工作负载上比 Ansor 快 6.1-9.5 倍。图 7（b）测量了 Conv2D 的速度。从 ResNet-50 中提取 batch size 为 32 的工作负载。所有 Conv2D 都使用 (3, 3) 卷积核和 (1, 1) 零填充。在所有情况下，BOLT 都比 Ansor 快 2.7-3.5 倍。总体而言，BOLT 实现了更高的性能。

图7 GEMM/Conv 的微基准测试

4.2 端到端网络的基准测试

通过对六个常见的卷积神经网络进行基准测试来评估 BOLT 在端到端模型优化方面的性能。以TVM Ansor 作为基准，按照官方示例进行配置，并将调优试验设置为 900 × 任务数。所有模型的 batch size 为32，数据类型为 FP16。推理速度如图 8（a）所示。可以看到，与 Ansor 相比，BOLT 的推理速度平均提高了 2.8 倍。在调优时间上，如图 8（b）所示，BOLT 可以在 20 分钟内完成调优，而 Ansor 平均需要 12 小时。

图8 常见卷积神经网络的基准测试

参考文献

• [1] L. Zheng, C. Jia, M. Sun, Z. Wu, C. H. Yu, A. Haj-Ali,Y. Wang, J. Yang, D. Zhuo, K. Sen et al., 14th USENIXSymposium on Operating Systems Design and Implementa-tion (OSDI 20), 2020, pp. 863–879.

• [2] J. Xing, L. Wang, S. Zhang, J. Chen, A. Chen and Y. Zhu,Bolt: Bridging the Gap between Auto-tuners and Hardware-native Performance, arXiv preprint arXiv:2110.15238, 2021.

• [3] Z. Chen, C. H. Yu, T. Morris, J. Tuyls, Y.-H. Lai, J. Roesch,E. Delaye, V. Sharma and Y. Wang, Bring Your Own Codegento Deep Learning Compiler, arXiv preprint arXiv:2105.03215,2021.

欢迎关注Adlik Github：https://github.com/Adlik/Adlik。