组合microkernel进行卷积算子优化

背景介绍

在广泛的科学和机器学习应用中，如何实现不同的优化张量计算的方法，在利用特定处理器结构的整体性能上是十分有必要的。同时，在接近处理器的最高性能时，微架构的复杂性使其变得非常困难。因此，提出关于张量计算的整体的主流的优化实现方案具有相当大的意义。当下，用于张量算子（针对 2D 卷积）的主流优化方案以及各自局限性包括：

• Polyhedral 编译器：该类方法可以为任何仿射循环计算（例如 2D 卷积）自动生成多级 tiled 代码。遗憾的是无法直接通过 tile 大小进行优化，而这是在 CNN 实现中十分必要的。

• Vender 库：该类方法由专家 HPC 和软件工程师进行手动优化，例如像 oneDNN 和 cuDNN。虽然这些实现使用 JIT 优化，但它们仍不能完全适应 DNN pipeline 中 CNN 层的每个张量尺寸。

• 自动调优：该类方法在确定优化空间后，可由系统通过机器学习算法自动完成从优化代码的生成，编译和执行，例如 AutoTVM 与 AutoScheduler。但目前尚未与高效优化的第三方库（像oneDNN）进行适配连接。

• 分析建模与优化：该类方法通过手动推理以构建数据移动和修剪的分析成本模型，并结合非线性优化问题的自动解决来优化 tile 大小。最内层的循环使用手动创建的微内核，类似于 BLIS one 和 Vender 库。实验显示该类方法同比优于其他 sota 方法，具有最优越的性能表现。但是以往的实现方法中往往底层基于单个微内核，仍需要部分 tiles 和 padding 进行辅助计算。

由 Tollenaere, Nicolas, et al 提出的 TTile，在充分考虑以上方案的优缺点的背景下，介绍了一种使用不同大小的 microkernel 组合 + 分析模型和优化的整体优化方案，以实现张量计算更加优越的加速效果。这是由于部分 tiles 和 padding 会造成额外的数据搬运和无效计算，而组合微内核可以消除该缺点，再加上分析建模与优化可以带来更一致与高性能的性能结果。

方法总览

首先从整体上进行方案介绍，其基本原理源于该领域内所近期取得的成果以及优化工作所面临的挑战：

• 将多级平铺循环计算精确到最内层的细粒度性能建模极具挑战性，并且对于像 CNN 这样的复杂计算，任何先前的工作都没有实现，但内存层次结构的粗粒度分析建模已被证明是有效的。他们通过使用固定的微内核使分析建模变得粗粒度且足够准确，但是这样一来他们限制了问题的搜索空间，很难保证同一 kernel 在不同问题上的一致性表现。

• 自动调整是克服精确性能建模挑战的有效方法。将其应用于多级平铺循环的主要障碍是优化空间的数量爆炸以及研究有效的遍历算法。AutoTVM 确实找到了很好的解决方案，但只能通过使用专家设计的脚本来限制搜索空间，往往需要设计者对特定硬件的内存管理等知识极其熟悉。

本文提出的方案通过使用一个两级分割的策略，在整个设计空间进行有效快速搜索：

• 通过广泛的一次性自动调整为给定的目标硬件平台开发一组微内核，这组微内核与要优化的 CNN 层的实际大小无关。然后，将它们组合起来，以完全适合问题大小，而不需要任何局部 tiles；

• 使用自动调整和分析建模的组合来显着修剪可能的 tile 循环排列和 degenerate 循环的空间。

Figure 1: 2D Convolution (unit stride)

Figure 2: Code generation sketch using microkernel composition. Convolution sizes are K = 64, C = 128,H = W = 136 and R = S = 1. Note that 136 = (8 + 13 × 2) × 4.

图 1 和 图 2 展示了对于卷积算子优化代码生成的完整方法。一个 2D 卷积包含 7 层嵌套循环，优化针对于不同层级的 tiling。对于一个 d 维的嵌套循环（这里 d=7）和一个 5 级内存分层（主存，L3~L1 caches，寄存器），各级 tiling 嵌套循环的总数为 5 d（此处为 35）。在图 2(a) 中显示为一组最外层的 d 块循环，这些循环遍历 L3 级块。每个 L3 级 tile 有 d 个 tile-loops 来遍历一组 L2 级 tile，依此类推，寄存器级 tile 被标记为微内核。

实际上，有效的 tiled 实现在每个存储层级只包含一个小的子集，即不需要进行 tile loops 的全部展开，有些循环在单次迭代中可以被优化，因此从代码中删除，但这些先验知识需要我们进行从内到外层的存储分析才能得到。于是首先确定单独运行良好的微内核。然后在其中挑选一两个适合所考虑问题大小的微内核进行组合。紧接着，使用分析建模来识别每个层级内的活动循环以及最终循环排列顺序。最后，使用自动调整来搜索外部 tiling 配置的显着修剪空间。图 2(b) 显示了该优化框架在一个用于目标平台的样本卷积上生成的代码，向量大小为 16 个元素。它使用两个微内核，一个对应于卷积迭代空间的切片，切片范围为 ，，，，另一个对应切片范围为 ，，，。这个例子说明了如何结合两个性能良好的微内核来划分问题大小。它还表明，在任何存储级别中只包含了部分 tile loops。因此证明了在使用自动调整进行探索之前，拥有一个分析模型来修剪搜索空间的重要性。

空间描述

迭代空间是指给定算子的循环索引所组成的整数向量集。优化策略 Tiling 是一种循环变换，它将迭代空间划分为集合，称为切片并以原子方式执行。这种划分允许我们控制每个 tile 访问的数据量，也就是 footprint，以确保它不超过给定的缓存容量。微内核是指由基本块组成的有效代码区域，该基本块由最内层并行循环的完全展开产生，包含在零个或多个嵌套的约减循环中。它通常用汇编语言或使用向量内在函数编写，旨在实现以下目标：

• 矢量 ALU 的有效利用；

• 通过展开和寄存器提升，跨迭代有效地重用（向量）寄存器中的数据；

• 足够的指令级并行性 (ILP) 以隐藏流水线功能单元的延迟（乘法和加法）。

由不同的优化策略共同组成最终的优化方案，从概念上讲，它可以被看作是一种专门的抽象，比 TVM 调度更高级别。优化方案是描述生成代码的分层结构的说明符列表，从最外层循环向内：

• 沿维度 插入外部循环。此循环将沿 遍历外层图块。这些 tiles 的大小应该恰好分割问题大小 。此外，对于给定的维度 ， 最多可能出现一次。

• 沿维度 插入一个 tile 循环。它沿着该维度精确地迭代 次。

• 用 虚拟插入一个 tile 循环，然后完全展开它（寄存器 tile）。

• 用 虚拟插入一个 tile 循环，其中 是向量长度，然后对其进行向量化。矢量化只发生在最内层：最多可能有一个 。

• , 是 的列表，沿维度 引入大小为 的 个循环序列。此符号代表生成组合嵌套的微内核 tile，构成最终的微内核。

通常来说，我们通过组合这些说明符可以用来描述针对优化空间的不同维度所进行的特定优化策略，最后对整个优化方案进行描述说明。

空间优化

TTILE 提出了一种新颖的优化算法，能够合理划分上一节中介绍的表达优化空间。该算法结合了分析建模和自动调整，它分为微内核的离线优化（寄存器级别），然后专门处理特定的卷积并决定平铺结构（缓存级别）。在 tiled 卷积的所有可能组合中，应用以下修剪策略：

• 只使用性能最好的微内核：这部分代码的性能限制了整个应用程序的性能，因此我们需要确保它尽可能高效地运行；

• 禁止部分 tiles，这会导致速度变慢，尤其是在生成代码的最内层；因为可分性约束对于某些问题规模来说太严格了，所以考虑微内核的组合。

整个算法总结如图 3：

Figure 3: Flow of the optimization algorithm.

1. 单独测量许多微内核的性能。微内核候选集由性能最佳（cpu使用率最高）的候选微内核组成。此阶段与问题大小无关，但特定于每个目标架构。

2. 对于每个候选微内核，考虑到问题大小和缓存大小，我们确定包含微内核的最佳循环排列（循环维度列表），即循环嵌套的主要结构（每个级别平铺的维度）。每个维度大小此时尚未确定。这种排列是通过当下分析跨越tiles下的 footprint、数据移动和重用而获得的。

3. 从选定的微内核和循环排列中，我们生成优化方案的空间。一般来说，我们要求在给定维度选择的任何 tile 大小是其子 tiles 大小的倍数，并沿该维度划分（完整）问题大小。除了微内核，这种限制是不可行的，举例来说：Yolo9000-8 有一个问题大小 ，这是质数，所有大小为 17 的微内核的性能都很糟糕。于是利用 说明符组成大小为 5 和 6 的微内核来解决问题：大小为 和 的微内核可以通过在前者的一次迭代之前对后者的两次迭代进行排序来组合，在不影响性能的情况下满足可分性约束。

4. 这时，搜索空间已经可以进行自动调整，但可以做得更好。通过引入一个简单的度量对所有生成的方案进行排序，我们可以从中检索任意数量的候选实现。推荐选择前 200 名。

5. TTILE 可以将表现最好的方案作为子树集成到 TVM 程序中。我们使用这种能力来促进性能比较，并自动将最优顺序方案升级为并行实现。

示例：考虑 AVX512 结构下微内核集合：，取 Yolo9000-13 ，问题大小为 （，循环排列确定为 。

维度 h 上的问题大小 H 为 34 = 2×17，因此所考虑的类中没有单个微内核与其一个除数相匹配。接下来，我们考虑该类的 2 个微内核的组合；2 × 11 + 12 就是这样一种组合。现在，我们需要将其他维度的数值分布在循环排列所描述的不同平铺级别上：

• k 维度是简单的：在微内核之上只有一个循环，并且沿着这个维度的微内核大小的占用空间为 32。因此，我们需要在外循环上平铺 16 倍以达到 512。

• c 维度也很简单：唯一的循环需要平铺 256 倍。同样，对于 r 和 s 维度，平铺系数都应该是 3。

• h 维度已经由微内核的组合管理。我们有 3 个位置可以放置 两个微内核的组成。我们考虑放在最外部。

• w 维度有 34 个循环分布在 2 个平铺级别。有 4 种组合：34 × 17 × 2、2 × 17 和 1 × 34。让我们考虑第二种。

于是优化后的调度方案为: 。

实验效果

TTILE 与几种 sota 方法进行了生成代码的性能评估实验：oneDNN（英特尔库，V2.3）、AutoTVM（自动调整，2021 年 6 月的 v0.8.dev0）、Autoscheduler（Ansor，自动调整 ，在 TVM 中）和 Mopt（分析建模）。设置：处理器型号（a）18 核 Intel Xeon Gold 5220 Cascade Lake 处理器（频率设置为 2.2 GHz）；（b）32 核 Intel Xeon Gold 6130 Skylake 处理器（频率设置为 2.3 GHz）。指令集均为AVX512，测试操作系统是 Debian GNU/Linux，内核版本 4.19，使用 PAPI 版本 5.7.0 监控硬件计数器。主要评估了两个网络的卷积性能：Yolo-9000 和 ResNet-18。它们的卷积层大小可以在图 4 中找到：

Figure 4: Convolution benchmarks and sizes. The kernels marked with a * are stride 2, else strideDimension k of Yolo9000-23 was padded to 28272 (a multiple of 16) to vectorize it on AVX512.

图 5 展示了性能结果，报告为多核 CPU 峰值性能的一部分。图中也展示了所有可用基准中每个工具和顺序/并行情况下的几何平均值。其中 MOpt 组机器性能标记为 0% 的数据是由于复现失败。我们注意到 TVM+Ttile 的平均水平远高于其他最先进的工具，这证明了该方法的可行性。尽管实验显示  oneDNN 的性能结果似乎低得惊人，但已和官方人员沟通确认。就找到最佳实现的总编译时间而言，在所有基准测试中，AutoTVM 和 AutoScheduler 都需要大约 5 小时。MOpt 大约需要 1h30，TTile 大约需要 2h。OneDNN 由于执行 JIT 编译，因此没有离线探索阶段。

Figure 5: Comparison with AutoTVM, AutoScheduler, oneDNN, Mopt for AVX512 (Intel Xeon Gold5220 and 6130), shown as percentage of machine peak. The left side of each line is the sequentialperformance and the right side is the parallel performance for the same convolution sizes.

在图 6 中报告了实验中的最佳性能方案使用的微内核。我们注意到23 种不同卷积操作一共使用了 16 种不同微内核。例如，为 ResNet18-1 选择的微内核的大小是 ，即沿 h 维度展开 14 次展开迭代，沿 k 维度进行 2 次迭代。相比之下，Yolo9000-4 使用了两个微内核 的组合。两个微内核沿 r 维度的大小为 3，但沿 h 维度的大小不同。第二个微内核的使用频率也是第一个微内核的两倍。

Figure 6: Microkernels of the best parallel scheme found by TTile on an Intel Xeon Gold 6130.

论文总结

为实现更高的推理性能，编译优化是必经之路，本篇介绍了 TTile，这是一种为张量计算生成高性能代码的工具，并在静态形状 2D 卷积上对其进行了评估。生成的代码利用大量高性能微内核并将它们组合起来以避免部分切片。该方法本身是不同类型的最先进方法的组合，包括分析建模、自动调整和实验。文内展示的性能结果在绝大多数情况下都证明了比 oneDNN 更高的性能，并且始终比最先进的 TVM 张量编译器和自动调优器以及 MOpt 具有更高的性能。更多相关内容知识，请参考引文。

参考文献

【1】Tollenaere, Nicolas, et al. "Efficient convolution optimisation by composing micro-kernels." (2021).

【2】Li, Rui, et al. "Analytical characterization and design space exploration for optimization of CNNs."Proceedings of the 26th ACM International Conference on Architectural Support for Programming Languages and Operating Systems. 2021.

【3】Olivry, Auguste, et al. "IOOpt: automatic derivation of I/O complexity bounds for affine programs."Proceedings of the 42nd ACM SIGPLAN International Conference on Programming Language Design and Implementation. 2021.

Adlik Github仓库