Roller: Fast and Efficient Tensor Compilation for Deep Learning

  • 文章标题: ROLLER：面向深度学习的快速高效张量编译
  • 作者/机构: Hongyu Zhu†⋄∗, Ruofan Wu‡⋄∗, Yijia Diao\$⋄∗, Shanbin Ke¶⋄∗, Haoyu Li§⋄∗, Chen Zhang£⋄∗, Jilong Xue⋄, Lingxiao Ma⋄, Yuqing Xia⋄, Wei Cui⋄, Fan Yang⋄, Mao Yang⋄, Lidong Zhou⋄, Asaf Cidon§, Gennady Pekhimenko†. (†多伦多大学, ‡中国人民大学, \$上海交通大学, ¶加州大学圣地亚哥分校, §哥伦比亚大学, £清华大学, ⋄微软研究院)

A1 主要贡献

本文介绍了ROLLER，一个深度学习张量编译器，旨在解决现有编译器（如TVM、Ansor）编译时间过长的问题。当前先进的编译器为深度神经网络（DNN）中的算子生成高效内核通常需要数小时，这极大地拖慢了DNN的开发周期。

核心问题: 现有DNN编译器将算子优化视为一个组合优化问题，需要在巨大的搜索空间（可达数百万种选择）中使用机器学习算法进行搜索。这个过程非常耗时，通常需要数千步搜索，每一步都需要在真实硬件上评估，导致编译一个模型可能需要数天甚至数周，尤其是在较新的加速器上。

研究目标: 提出一种全新的、基于构造而非搜索的方法来快速生成高效的算子内核，将编译时间从小时级降低到秒级，同时保持与最先进解决方案相当甚至更好的性能。

核心创新点:
1. 数据处理流水线视角: ROLLER不将算子计算视为多层嵌套循环，而是看作一个数据处理流水线。在这个流水线中，数据分片（tile）在具有并行执行单元和多层内存结构的抽象硬件中被移动和处理。优化目标转变为最大化该流水线的吞吐量。
2. rTile抽象: 提出了一个名为rTile的新抽象。rTile封装了与底层加速器关键特性（如内存bank、内存事务长度、最小可调度单元）对齐的张量形状。通过强制对齐，rTile极大地限制了可选的tile形状，使得构造高效流水线变得可行且高效。
3. 递归构造算法: 采用一种“先扩展后推广”（scale-up-then-scale-out）的递归构造算法。首先，通过逐步增大rTile的尺寸，构造出一个能使单个执行单元（如GPU的SM）饱和的程序（rProgram）。然后，利用DNN计算和硬件的同质性，将该程序复制到所有并行执行单元。
4. 微观性能模型: 由于rTile与硬件的对齐特性，其性能（如内存吞吐量）是高度可预测的，可以通过硬件规格或微基准测试分析得出。这使得ROLLER能用一个轻量级的微观性能模型高效评估不同rTile配置，避免了现有编译器所需的昂贵的在线性能评测，从而极大地加速了编译过程。

通过这些创新，ROLLER能够在数秒内生成与最先进编译器和供应商库性能相当甚至更优的内核，编译速度提升了三个数量级，并且能轻松适配不同的加速器，如NVIDIA GPU、AMD GPU和Graphcore IPU。

图1：不同tile形状的访问模式。矩阵乘法，Cm,n = Am,k × Bk,n。
图1：不同tile形状的访问模式。矩阵乘法，Cm,n = Am,k × Bk,n。

A3 背景知识/关键Observation

过长的编译时间。我们的经验表明，使用最先进的张量编译器Ansor【33, Ansor: Generating high-performance tensor programs for deep learning+2020+OSDI 20】，编译单个算子的平均时间为0.65小时，其中ResNet模型中的一个卷积算子耗时2.17小时。一个DNN模型可能包含数百个算子，编译整个模型轻易就需要数天。例如，编译一个NASNet模型【36, Learning transferable architectures for scalable image recognition+2018+CVPR】，在调优41.8小时后，我们仅完成了32%的搜索进度。在不够成熟的设备上，编译速度甚至更慢。

观察与洞见。我们观察到，对DNN算子的计算存在一种不同的视角。以矩阵乘法（MatMul），$C_{m,n} = A_{m,k} \times B_{k,n}$ 为例，现有编译器将其视为一个遍历m、k、n三个轴的三层循环，而这个计算过程也可以被看作一个数据处理流水线：从A和B中加载（Load）子矩阵（即tile），计算（Compute）这两个tile，并将结果tile存储（Store）到C的内存中。因此，计算性能取决于在这个Load-Compute-Store流水线中移动数据tile的速度。影响流水线所有步骤性能的关键因素是tile的形状及其在一维内存空间中的布局。图1(a)展示了计算C中一个元素时的内存访问模式。假设所有矩阵按行主序存储，加载B的一列会导致步长为1的跨步访问。如果内存事务长度为4，将有3/4的数据读取是冗余的。因此，数据tile的形状应与内存事务长度对齐以实现高效内存访问。在图1(b)中，当以1×4的tile粒度计算B时，就不会有内存带宽浪费。除了内存对齐，tile形状还应与硬件执行单元（如并行线程数）对齐，以避免计算周期的浪费。此外，tile形状还影响缓存带来的数据复用机会。例如，图1(a)在每次计算1×1的tile时需要2mnk次数据读取。而在图1(b)中，由于A的一次读取可以被复用4次，仅需1.25mnk次读取。如果将M维度的tile大小设为4×4，如图1(c)所示，总读取次数可以减少到0.5mnk，相比图1(a)提升了10倍。这些观察启发了ROLLER，一个旨在识别对齐的tile形状并构建高效tile处理流水线以提高端到端吞吐量的系统。

图2：ROLLER系统概览。
图2：ROLLER系统概览。

A2 方法细节

3. 系统设计

ROLLER系统概览。如图2所示，ROLLER接收一个以张量表达式描述的算子。该表达式由用户生成或来自图级DNN编译器【15, TVM: An automated endto-end optimizing compiler for deep learning+2018+OSDI 18】【26, Rammer: Enabling holistic deep learning compiler optimizations with rtasks+2020+OSDI 20】【33, Ansor: Generating high-performance tensor programs for deep learning+2020+OSDI 20】，后者可能已将多个算子融合成一个表达式。ROLLER从张量表达式中提取张量形状，并利用硬件规格来构建rTile，这是一个与硬件对齐的构建块（§3.1）。基于rTile，ROLLER提出了一种“先扩展后推广”的递归构造算法，用于生成描述数据处理流水线的高效张量程序（称为rProgram）（§3.2）。在生成rProgram时，构造算法通过一个微观性能模型评估rProgram的性能来识别好的rTile配置。该模型建立在一个硬件抽象层之上，该抽象层仅暴露与rTile相关的接口：Load、Compute和Store（§3.3）。构造出的rProgram最终通过代码生成器为特定设备生成最终的内核代码。

图3：rTile的数据结构。
图3：rTile的数据结构。
图4：由一个rTile为MatMul表达式推断出的数据tile和计算tile。
图4：由一个rTile为MatMul表达式推断出的数据tile和计算tile。

3.1 张量表达式与rTile

rTile的定义与推导。ROLLER接收的张量计算输入是一种基于索引的lambda表达式，即张量表达式【15, TVM: An automated endto-end optimizing compiler for deep learning+2018+OSDI 18】【27, Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines+2013+PLDI ’13】。它描述了输出张量中的每个元素如何根据输入张量中相应元素计算得出。例如，一个输出张量C形状为M×N的MatMul算子可以表示为 $C[i, j] = \sum_k A[i, k] \times B[k, j]$，其中C中索引为(i, j)的元素是通过对A的第i行和B的第j列的元素进行求和规约计算得到的，k是规约轴。这种表达式能覆盖DNN模型中的绝大多数算子，并被现有DNN编译器广泛使用。ROLLER引入了RollingTile（简称rTile）作为构成张量计算的基本计算单元。如图3所示，一个rTile封装了一个沿给定张量表达式expr的每个循环轴定义的多维tile形状。给定shape和expr，rTile可以静态推断出所涉及的输入和输出数据tile。例如，对于上述MatMul表达式，一个沿i、j、k轴的tile形状[4, 4, 2]表示一个rTile，其中每个rTile从A加载一个4×2的数据tile，从B加载一个2×4的tile，进行总共4×4×2次乘加计算，并将一个4×4的数据tile存储到C，如图4所示。

图5：(a) 事务对齐的内存加载和 (b) 无bank冲突的填充示意图。
图5：(a) 事务对齐的内存加载和 (b) 无bank冲突的填充示意图。

rTile的对齐属性。rTile的一个独有属性是它必须与底层硬件特性和给定张量表达式中的张量形状都对齐。所有这些对齐都由rTile的shape和storage_padding字段（如图3所示）控制，它们分别代表了rTile的逻辑形式和物理布局。我们接下来详细阐述对齐的具体要求。

与硬件执行单元对齐。首先，rTile的形状必须与其运行的执行单元的并行度对齐。例如，如果在GPU的一个warp线程上运行，rTile的shape大小应为warp大小（如32）的倍数，以实现最大计算效率。当使用NVIDIA GPU的TensorCore时，rTile的形状应为16×16×16的倍数。类似地，在流式多处理器（SM）上执行的rTile应将其大小与SM上的执行单元数量对齐。

与内存事务对齐。其次，数据tile的形状应与内存事务的长度对齐，以实现最优内存访问。具体来说，对于rTile的每个数据tile，我们应保证其主导维度（例如，行主序张量中最内层维度）是内存事务长度的倍数，如图5(a)所示。在ROLLER中，张量以缓存对齐的方式分配。因此，一个rTile可以避免任何事务读取的浪费，因为其形状与内存事务对齐。

与内存bank对齐。第三，数据tile的内存布局应使其步长与内存bank对齐，以避免读取冲突。例如，一个[3, 4]的数据tile存储在跨4个bank的内存中，并由一个形状为[3, 1]的上层内存tile读取，如图5(b)所示。将所有[3, 1]的值存储在同一个bank中的朴素方法会导致加载冲突。rTile通过填充数据tile来避免这种低效。给定一个主导维度大小为N的数据tile，它被另一个主导维度大小为n的tile读取，我们在存储该tile时沿N维度添加一个大小为 $(B \times L - N\%(B \times L) + L \times \lceil n/L \rceil) \% (B \times L)$ 的填充，其中B和L分别是bank数量和bank宽度。沿每个轴的填充大小被计算并存储在图3的storage_padding字段中。对于图5(b)中的情况，通过大小为1的填充，所有[3, 1]的值被分布在不同的bank中，可以被高效地读取。

与张量形状对齐。最后，rTile的形状应与输入张量表达式的张量形状对齐。否则，计算无法被rTile均匀划分，会浪费计算资源或产生沉重的边界检查开销。一个简单的解决方案是沿大小为$N_i$的张量维度i添加一个填充大小$P_i$，使得$N_i + P_i$成为rTile形状在该轴上维度大小的倍数。然而，大的填充可能会浪费计算。因此，ROLLER将张量填充限制在一个范围$\epsilon$内，其中rTile的形状维度大小$S_i$必须满足 $(S_i - N_i \% S_i) / N_i \le \epsilon$，这里$N_i$是维度i的张量大小。这确保了计算的浪费百分比受$\epsilon$限制。有了这个限制，我们可以枚举所有满足此条件的有效rTile形状。

推导所有rTile。基于上述对齐要求，对于一个特定的张量表达式和硬件设备，ROLLER使用以下接口增量地推导所有符合条件的rTile：vector<int> GetNextAlignedAxisSize(rTile T, Dev d)，该接口返回给定rTile T和特定设备规格d时，其形状中每个轴的下一个对齐尺寸。这是通过沿每个轴逐渐增加维度大小直到满足所有对齐要求来计算的。rTile抽象允许ROLLER通过向GetNextAlignedAxisSize接口添加新的对齐规则来扩展以支持新的对齐要求（例如，新的硬件特性）。

计算数据复用分数。rTile的一个有趣特性是我们可以通过调整其形状来隐式控制内存流量。增加rTile大小通常会以占用更多内存空间为代价带来更多的数据复用机会。给定一个rTile T及其在每个轴上的下一个对齐尺寸，我们可以计算轴i的数据复用分数$S_i = \frac{Q(T) - Q(T'_i)}{F(T'_i) - F(T)}$，其中$T'_i$是一个新的、通过将轴i的维度大小替换为GetNextAlignedAxisSize返回的下一个对齐尺寸而扩大的rTile。函数$Q(T)$和$F(T)$计算当计算以T为粒度执行时的内存流量和内存占用，这可以直接根据给定的张量表达式和硬件内存规格推断出来（§3.3）。更大的$S_i$意味着更好的成本效益，即用相同的内存使用量可以节省更多的内存流量。内存复用分数在构建高效rProgram（使用rTile）中扮演着关键角色，如下一小节所示。

3.2 张量程序构造

rTile程序 (rProgram)。鉴于rTile和现代加速器的分层内存结构，张量计算可以很自然地被视为一个流式数据处理流水线。计算过程将数据tile（在rTile中指定）从最低层内存通过内存层次结构加载到最高层，在加速器的执行单元上执行rTile计算，然后将结果数据tile存回最低层内存。对于每个内存层，都会定义一个特定的rTile以与该内存层的特性对齐。因此，ROLLER将张量计算描述为一个具有分层rTile配置的数据处理流水线，称之为rTile程序（即rProgram）。

图6：一个在具有3层内存层级（自底向上：L2层到L0层）的设备上的rProgram伪代码。
图6：一个在具有3层内存层级（自底向上：L2层到L0层）的设备上的rProgram伪代码。
图7：ROLLER计算模型。(a)一个rTile程序；(b)矩阵乘法上的rTiles；(c)rTile程序在硬件内存层次结构上的执行。
图7：ROLLER计算模型。(a)一个rTile程序；(b)矩阵乘法上的rTiles；(c)rTile程序在硬件内存层次结构上的执行。

rProgram示例。图6展示了一个在具有三层内存（L0、L1和L2）的设备上的rProgram。该rProgram由每层的rTile以及每层的rTile指令（即Load、Store和Compute）描述。图7(a)展示了一个在图7(b)中图示的MatMul rProgram。图7(c)说明了该rProgram如何映射到设备的每个内存层。具体来说，每次它从内存L2加载一个[4, 4]的数据tile到A和一个[4, 8]的tile到B至L1；然后它以[2, 1]和[1, 2]的形状将数据tile从内存L1加载到内存L0（即寄存器）。每次Compute之后，产生的[2, 2] tile将直接从L0存储到L2。

rProgram的构造策略。对于一个数据处理流水线，相应的rProgram的优化目标是最大化流水线的吞吐量。这个目标可以转化为三个条件：1）计算和内存移动应充分利用硬件特性；2）吞吐量应使瓶颈阶段饱和；3）需要有足够的并行性以利用所有并行执行单元。因此，ROLLER提出了以下的rProgram构造策略：首先通过构建一个单核rProgram来在单个核心上进行扩展（scale-up），以使核心的硬件利用率饱和，然后通过复制构造出的单核rProgram来推广（scale-out）以利用多核并行性。

图8：ROLLER针对单个核心（如一个SM）的rProgram构造算法。
图8：ROLLER针对单个核心（如一个SM）的rProgram构造算法。

扩展一个rProgram。由于rTile的对齐属性保证了硬件效率，ROLLER只需专注于通过构造正确的rTile形状来最大化每个内存层的吞吐量。通过利用§3.1中定义的数据复用分数，单核rProgram构造算法从一个初始的rTile开始，并逐渐朝着rTile中最具成本效益的轴（即具有最大数据复用分数的轴）扩大它。请注意，构造算法不需要绝对的数据复用分数，它只是选择最大的一个来最大化吞吐量。在此过程中，内存性能不断提升，直到达到计算瓶颈或最大内存容量。上述过程从顶层到底层为每个内存层重复进行，直到构造出一个理想的rProgram。注意，如果某些张量表达式（如逐元素算子）的数据复用分数保持不变，ROLLER将只为顶层构造rTile并直接从底层内存加载它们。图8展示了详细的构造算法。给定一个张量表达式expr和一个目标设备dev，算法在顶层内存构造一个初始rTile T，并递归地扩大T（第4行的EnlargeTile）。在每一步，它枚举下一个能最大程度提高数据复用分数的更大rTile $T'$（第10行的GetNextRTileShapes）。如果$T'$达到内存容量（第13行）或数据tile加载吞吐量$MemPerf(T')$超过峰值计算吞吐量$MaxComputePerf(T')$（第17行），算法会记录当前的rTile T，并继续在下一个内存层进行EnlargeTile。否则，它会继续在当前层扩大$T'$（第20行）。构造在最底层内存结束（第6行），产生一个结果并重复，直到获得K个（例如5-20个）rProgram（以容忍设备编译器带来的隐藏因素影响）。注意，$MemPerf(T')$和$MaxComputePerf(T')$是基于dev，根据微观性能模型（§3.3）推导出来的。

推广一个rProgram。鉴于大多数DNN算子的计算模式和加速器中并行执行单元的同质性，ROLLER通过将计算均匀地划分为大小等于最底层rTile的tile，简单地将一个执行单元上构造的rProgram复制到其他单元。我们通过将所有分区均匀地分配给所有执行单元来实现这一点。请注意，ROLLER倾向于将沿规约轴拆分的分区分配给同一个执行单元，因为它们可以在更高层的内存中共享规约结果。另外，ROLLER不假设一个rProgram会独占所有计算单元，系统可以在推广时明确控制rProgram的并行度。

小算子和不规则张量形状处理。推广算法天然地偏好具有充足并行性的算子，例如分区数量远大于执行单元数量的情况。对于小型算子，算法的整体性能可能会因并行执行单元利用率低而受损。通常，如果存在足够的算子间并行性，这可以通过在像Rammer【26, Rammer: Enabling holistic deep learning compiler optimizations with rtasks+2020+OSDI 20】这样的编译器中进行协同调度来解决。否则，对于每个rProgram，ROLLER将尝试沿数据复用分数最小的轴缩小其rTile以实现足够的并行性。注意，这个枚举过程每次都像其他对齐规则一样返回下一个对齐的tile大小，这是一个高效的过程，与整体构造过程相比，其成本可以忽略不计。此外，一个大算子可能包含具有小维度的不规则张量形状，ROLLER可能因对齐要求而无法生成足够数量的rProgram。为解决此问题，ROLLER通过一个轴融合（axis fusion）过程将张量表达式转换为规范形式。具体来说，对于所有涉及的张量，如果一个张量中存在两个相邻的轴，这两个轴在所有其他张量中要么都存在且仍然相邻，要么都缺失，ROLLER可以安全地合并这两个轴。例如，对于一个输入和输出张量形状均为[17, 11, 3]的逐元素算子，ROLLER会通过融合三个轴将其转换为形状[561]（17×11×3）的张量。除了轴融合，ROLLER还会尝试贪婪地增加张量填充机制中的参数$\epsilon$（§3.1），直到构造出K个rProgram。

图9：ROLLER的硬件抽象接口。
图9：ROLLER的硬件抽象接口。

3.3 高效评估一个rProgram

评估机制。在构造算法中，ROLLER需要评估rProgram的性能。ROLLER不是在真实硬件设备上评估端到端的rProgram，而是只需要评估相应rTile的性能，例如图8中的MemPerf和MaxComputePerf。为此，ROLLER构建了一个针对硬件抽象层（HAL）描述的设备的微观性能模型。HAL将加速器建模为多个具有分层内存的并行执行单元。HAL暴露了三个基于rTile的接口：Load、Compute和Store（图9）。一个执行单元被抽象为一个rTile执行单元（TEU），它通过Compute接口计算数据tile。多个TEU可以组织成一个组，协同地Load和Store tile。HAL将不同的内存层，如寄存器、共享内存、DRAM，视为一个统一的类型，暴露影响tile移动性能的硬件规格。这些规格包括内存容量、事务长度、缓存行大小和内存bank数量，这些都可以通过图9中的GetDeviceSpec接口获得。

微观性能模型。借助硬件抽象层，ROLLER可以轻松推导出rTile（及其rProgram）的性能。首先，给定一个rTile，其产生的内存占用（包括填充）和跨不同层的内存流量可以从rTile的张量表达式expr和shape静态推断出来，即图9中的MemFootprint和MemTraffic接口。它们用于计算数据复用分数并检查rTile是否超过内存容量。其次，为了计算rTile的MaxComputePerf，ROLLER进行一次性性能分析，通过积极扩大计算tile（例如，SM中warp大小的倍数）来测量峰值计算吞吐量以使TEU饱和。此性能数据在ROLLER中被缓存，以供构造算法未来查询。最后，对于给定的rTile，ROLLER还估计MemPerf，即从一个内存层向更高层加载数据tile的性能。鉴于rTile中的内存访问是对齐的，加载一个规则数据块的延迟可以简单地通过总流量除以内存带宽来建模。对于所有TEU共享的内存层，我们平均分配带宽。对于较小的访问大小，ROLLER也为每种设备类型进行一次性离线性能分析并缓存结果。值得注意的是，微观性能模型仅需在tile形状完全对齐时保持准确，这是ROLLER的一个关键要求。

4. 实现

实现基础。我们的ROLLER实现基于TVM【15, TVM: An automated endto-end optimizing compiler for deep learning+2018+OSDI 18】和Rammer【26, Rammer: Enabling holistic deep learning compiler optimizations with rtasks+2020+OSDI 20】这两个开源DNN编译器。ROLLER的核心机制，包括表达式优化、构造算法、微观性能模型等，用8000行代码实现。ROLLER的编译流程如下：输入是一个ONNX图【9, ONNX. https://onnx.ai/】或TensorFlow冻结图 【13, TensorFlow: A System for Large-Scale Machine Learning+2016+OSDI 16】。ROLLER首先利用Rammer进行图级优化（例如，算子间和算子内的协同调度）。接下来，ROLLER为从优化图中提取的每个（融合的）算子推导TVM张量表达式，并通过ROLLER的构造算法生成相应的rProgram，并进行内核生成。最后，生成的内核被注入到Rammer的运行时，生成端到端的模型代码。

代码生成。鉴于rProgram中固定的代码结构（如图6所示），ROLLER通过一个预定义的模板生成内核代码，该模板作为TVM调度及其内置调度原语实现。在每个内存层加载和存储数据tile是通过TVM的cache_read和cache_write原语实现的。对rTile的分区是通过split和fuse完成的。一些基本的rTile计算是用TVM的内部API实现的。有了这个模板，一个给定的rProgram可以直接生成设备代码，例如CUDA内核。

张量填充。ROLLER依赖张量填充来使rTile与张量形状对齐。在实践中，最低层内存（如DRAM）中的大多数张量是由外部程序（如DNN框架）分配的，因此我们只在上层内存（如共享内存）中应用填充。我们的张量填充目前要求输入张量表达式指定是否允许填充，以及默认的填充值（例如，MatMul算子为0）。对于内存bank对齐的存储填充，我们利用TVM的storage_align原语添加填充。

性能分析。ROLLER实现了两个分析器：一个微观性能分析器和一个内核分析器。前者通过一组微基准测试生成设备规格，例如内存带宽、计算吞吐量等，这是对每种设备类型和张量表达式类型（不考虑张量形状）的一次性离线分析。后者分析前K个rProgram中最快的内核，并且如果K大于1，则用于每个编译结果。在实践中，特定内核代码的性能也受到一些设备编译器和硬件相关的隐藏因素的轻微影响，ROLLER很难控制这些因素，包括不同指令类型的指令密度、寄存器分配行为、设备编译器优化、warp调度开销等。特别是在NVIDIA GPU上，ROLLER依赖nvcc【3, CUDA NVCC. https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/】将生成的CUDA代码编译成机器码。然而，nvcc的专有优化可能会对程序执行行为产生不良影响。因此，ROLLER利用内核分析器快速评估性能最好 的rProgram并选择最佳的一个。较大的K通常可以提高内核质量。在评估了前10、20和50个结果后，我们的经验表明，前10个结果在大多数情况下可以找到最优结果。注意，ROLLER的内核分析器与之前编译器中由机器学习算法驱动的评估过程不同【15, TVM: An automated endto-end optimizing compiler for deep learning+2018+OSDI 18】【33, Ansor: Generating high-performance tensor programs for deep learning+2020+OSDI 20】【35, Flextensor: An automatic schedule exploration and optimization framework for tensor computation on heterogeneous system+2020】。基于ML的方法通常需要数百甚至数千个顺序评估步骤，而ROLLER只并行分析几十个候选者。未来，我们计划实现汇编级代码生成，以减轻高级设备编译器中的隐藏问题。

多加速器支持 (HAL)。ROLLER的HAL使我们能够轻松支持不同的加速器。用户可以为每种设备类型配置相应的HAL。ROLLER也为大多数常见设备类型提供了内置配置。一些详细配置，例如内存带宽，依赖于微基准测试分析或从已发布的设备规格中推导。接下来，我们分享在几种流行的DNN加速器上实现HAL的经验，包括NVIDIA GPU、AMD GPU和Graphcore IPU。

在NVIDIA CUDA GPU上实现ROLLER。NVIDIA GPU通常采用集中式内存架构。我们在V100和K80这两种具有不同流式多处理器（SM）架构的CUDA GPU上实现了ROLLER。它们的内存架构包含全局内存、L2缓存、L1缓存、共享内存和寄存器。在ROLLER的HAL中，我们将它们抽象为3个内存层：L2层用于全局内存和L2缓存，L1层仅用于共享内存，L0层用于寄存器。我们忽略L1缓存，因为它与共享内存共享空间，并且不能被用户程序控制。所有级别的内存带宽都由我们的微基准测试测量。全局内存层的事务长度对两种GPU都设置为32字节，即8个浮点元素。对于V100 GPU，共享内存的bank数量和bank长度分别为32和4字节。对于K80 GPU，bank长度为8字节。两种GPU的共享内存容量都设置为48KB（基于deviceQuery）。我们将CUDA GPU上的TEU实现为一个32线程的warp，这也是执行TensorCore WMMA指令的基本单元。一个HAL上的TEU组的大小（例如一个SM）被设置为warp调度器的数量，两种GPU都为4。V100的SM数量为80【21, Dissecting the nvidia volta gpu architecture via microbenchmarking+2018+arXiv】，K80为13。在CUDA GPU上，每个线程的寄存器容量有限，例如V100为255个寄存器。超过此限制将导致寄存器溢出，造成显著的性能下降。这为寄存器层的rTile大小设置了限制。我们注意到nvcc编译器会隐式声明更多寄存器（用于循环变量或其他目的）。鉴于这种行为难以预测，我们凭经验将V100和K80每个线程的寄存器限制减少到仅96个，以避免意外的性能影响。

在AMD ROCm GPU上实现ROLLER。我们还在AMD的第二代Vega系列GPU MI50【12, AMD Radeon Instinct™ MI50 Accelerator, accessed 2018 Nov】上实现了ROLLER。MI50的内存架构与V100相似：所有计算单元（CU）都可以访问集中的全局内存。与NVIDIA GPU中的SM类似，每个CU都有自己的暂存内存、寄存器和计算核心。ROCm【1, AMD ROCm Platform. https://github.com/RadeonOpenCompute/ROCm】内核程序的数据移动也类似。全局内存的内存事务大小设置为64字节。内存bank数量为32，bank长度也为4字节。我们也将TEU实现为一个线程warp， 在MI50 GPU上为64个线程。最大寄存器大小凭经验限制为每个线程70个。所有其他规格，如每层的内存带宽、峰值计算吞吐量等，都用我们的微基准测试测量。

在Graphcore IPU上实现ROLLER。Graphcore IPU【22, Dissecting the graphcore ipu architecture via microbenchmarking+2019+arXiv】是一个拥有1216个并行处理核心的大规模并行MIMD处理器。与NVIDIA和AMD GPU不同，IPU采用分布式内存架构。每个核心只有256KB的片上本地内存，没有统一的全局内存。当本地内存无法容纳所有输入数据时，默认情况下，内核程序的初始数据被存放在片上本地内存中，并均匀分布在各个节点上。因此，ROLLER的IPU HAL也抽象了三个内存层：L2用于所有核心的远程内存，L1用于每个核心的本地内存，L0用于寄存器。我们利用了先前的基准测试工作【22, Dissecting the graphcore ipu architecture via microbenchmarking+2019+arXiv】，该工作已成功测量了峰值内存带宽和计算吞吐量。每个IPU核心的寄存器文件大小未公开。考虑到我们无法预测IPU程序编译器的行为，我们允许每个上层rTile仅使用10个寄存器，这可以安全地保证tiling算法不会发出无效的tiling配置。

A4 实验环境

  • 数据集/模型:

    • 使用了四个典型的DNN模型进行基准测试：ResNet-50 (CNN), LSTM (RNN), NASNet (通过神经架构搜索获得的CNN模型), และ BERT-Large (基于Transformer)。
    • 默认批处理大小设置为128。
    • 从每个模型中选取最常用的算子构建了一个包含6类共119个不同配置实例的算子基准库（7个MatMul, 44个Conv2D, 23个DepthwiseConv, 28个逐元素算子, 13个池化算子, 4个规约算子）。

  • 硬件配置:

    • NVIDIA GPU:
      • Azure NC24s_v3虚拟机，配备Intel Xeon E5-2690v4 CPU和4块NVIDIA Tesla V100 (16GB) GPU。
      • Azure NC24_v1虚拟机，配备Intel Xeon E5-2690v3 CPU和4块NVIDIA Tesla K80 GPU。
    • AMD GPU:
      • 一台服务器，配备Intel Xeon E5-2640 v4 CPU和4块AMD Radeon Instinct MI50 (16GB) GPU。
    • Graphcore IPU:
      • Azure ND40s_v3虚拟机，配备Intel Xeon Platinum 8168 CPU和16个IPU。

  • 软件配置:

    • NVIDIA GPU平台: Ubuntu 16.04, CUDA 10.2, cuDNN 7.6.5。
    • AMD GPU平台: Ubuntu 18.04, ROCm 4.0.1。
    • Graphcore IPU平台: Poplar-sdk 1.0。
    • 对比系统: TVM (v0.8), Ansor (v0.8), cuDNN, cuBLAS, rocBLAS, POPLAR库, TensorFlow (v1.15), TensorFlow-XLA, TensorRT (v7.0)。
    • 实现基础: ROLLER基于TVM和Rammer实现。

表1: 基准测试中部分算子配置。
表1：基准测试中部分算子配置。

A4 实验结果

5.1 在NVIDIA GPU上的评估

  • 算子性能 (V100):

    • 与CUDA库(CudaLib)对比: ROLLER为81.5%的算子生成了性能相当（10%以内）的内核，并在59.7%的算子上甚至更快。性能较差的主要是3x3及更大滤波器的卷积，因为cuDNN使用了更高效的Winograd算法，而这难以用张量表达式表示。
    • 与TVM和Ansor对比: ROLLER分别为72.3%和80.7%的算子生成了性能相当的内核。在54.6%和65.5%的算子上，ROLLER生成的内核比TVM和Ansor更快，尤其是在执行时间超过5ms的大型算子上，平均加速比分别为1.85倍和1.27倍。对于性能较差的小型或不规则算子，其平均内核时间本身就很短（约1-1.65ms）。
      图10：V100 GPU上的算子性能（y轴：平均内核执行时间，单位ms）。

  • 编译时间 (V100):

    • TVM和Ansor的平均编译时间分别为0.65小时和0.66小时，最长可达7.89小时和2.17小时。
    • ROLLER生成top-1内核平均耗时仅0.43秒，快了三个数量级以上。即使是评估top-10候选内核，平均编译时间也仅为13.3秒。
      图11：每个算子的编译时间。

  • 可扩展性 (V100):

    • 随着算子尺寸（批大小）增加，ROLLER和Ansor在MatMul上表现出线性扩展性，性能与cuBLAS相当（图12）。
    • 在Conv2D上，ROLLER同样线性扩展，性能优于Ansor和TVM，平均分别快1.25倍和1.54倍（图13）。
    • 编译时间方面，TVM和Ansor的编译时间随尺寸增长而显著增加，而ROLLER保持稳定，top-1内核在1秒内，top-10在16秒左右（图14）。
      图12：BERT-Large模型中不同M尺寸的MatMul算子内核时间，K=1024, N=4096。
      图13：不同批大小N的Conv2d算子内核时间，C=1024, H=14, F=2048, K=1, S=2。
      图14：不同批大小下MatMul和Conv2d算子的编译时间。

  • TensorCore性能 (V100):

    • ROLLER通过对齐rTile形状，能快速为TensorCore生成良好内核，性能与高度优化的cuBLAS差距在43%以内。相比之下，TVM在4个测试算子中有3个无法生成有效内核。
      图15：TensorCore上的Matmul内核时间。

  • 小算子与不规则形状处理:

    • 对于并行度不足的小算子，通过收缩rTile（Roller-S）可比原始配置（Roller-O）平均提升2.3倍性能，但仍比Ansor慢50%（图16）。
    • 对于不规则形状，轴融合（Roller-F）比基线（Roller-B）平均提升1.5倍性能。进一步使用张量填充（Roller-P1.0）可再提升1.4倍性能（图17）。
      图16：小算子的性能。
      图17：不规则形状算子的性能。

  • 微观性能模型验证:

    • 实验表明，模型在配置与硬件并行度对齐时（如线程数>=128，块数>=80）能准确预测吞吐量，这证明了ROLLER只关注对齐配置的策略是合理的（图18）。
    • 大部分（66%）生成的内核能利用超过90%的硬件理论资源极限，证明了其高质量（表2）。
    • 不同的对齐优化（内存事务、执行单元、形状、bank）累积带来了1.94倍的性能提升（表3）。
      图18：微观性能模型的内存吞吐量（DRAM和共享内存）和计算吞吐量与实际测量的对比。
      表2：不同内核在饱和层的资源利用率分布。
      表3：不同对齐优化的平均累积性能提升。

  • 端到端模型性能 (V100):

    • 在BERT-Large和LSTM模型上，ROLLER（Rammer+Roller）性能超越了所有对比系统，分别比最先进的TF快1.07倍，比TensorRT快1.55倍。
    • 在ResNet和NASNet上，由于卷积核效率低于cuDNN，性能略逊于TF、TF-XLA和TF-TRT。
    • Ansor虽然性能有竞争力，但编译时间极长（平均29.3小时），而ROLLER编译整个模型平均只需422秒。
      表4：V100 GPU上的端到端模型执行时间（毫秒）和编译时间。

  • K80 GPU性能:

    • ROLLER在K80上为82.4%的算子生成了比CUDA库更好的内核。与TVM和Ansor相比，也分别在65.5%和71.4%的算子上表现更优。
    • 编译时间方面，ROLLER top-1内核平均5.24毫秒，top-10为12.3秒，远快于TVM（0.65小时）和Ansor（0.95小时）。
      表5：NVIDIA K80 GPU上性能更好和相当的算子百分比。

5.2 在其他加速器上的评估

  • AMD ROCm GPU (MI50):

    • ROLLER为73.1%的算子生成了比ROCm库更优的内核，显示出在不够成熟的生态中优势更明显。
    • 相比TVM和Ansor，也分别在58.8%和70.6%的算子上胜出。
    • 编译时间优势依然巨大：ROLLER top-10平均7.69秒，而TVM和Ansor平均需要0.85小时和0.99小时。
      表6：AMD ROCm MI50 GPU上性能更好和相当的算子百分比。

  • Graphcore IPU:

    • 在IPU上，ROLLER生成的内核全面优于Graphcore官方的Poplar-sdk库，平均快3.1倍，最高达9.2倍。
    • 与Ansor相比，ROLLER在大多数算子上性能相当或更好，平均提升2.9%。
    • 同样，ROLLER在几分钟内即可完成编译，而Ansor需要数小时，这在设备编译器本身较慢的新硬件上尤为重要。
      图19：Graphcore IPU上的算子性能（y轴为对数刻度）。

A7 补充细节

6. 讨论与未来工作

与基于循环的编译器的优化空间对比。rTile和数据处理流水线的抽象使得ROLLER的优化空间与现有DNN编译器（如Ansor【33, Ansor: Generating high-performance tensor programs for deep learning+2020+OSDI 20】）重叠但又有所不同。这些编译器将张量编译视为嵌套循环优化。例如，Ansor只允许沿张量维度的可整除的tiling大小来均匀划分循环轴，这使得它在处理具有素数维度的张量形状时通常表现较差。ROLLER则从数据处理流水线的角度关注最大化硬件效率，允许更激进的优化，例如探索不可整除但与硬件对齐的tiling大小，并结合融合相邻轴和填充张量形状。基于我们观察到大多数DNN算子是内存密集型的，ROLLER通过首先优化数据tile吞吐量（即最大化复用分数奖励和与硬件特性对齐），然后才考虑并行性，这与现有DNN编译器根本不同。这种权衡固有地导致了快速编译和对具有足够并行性的算子的高性能。

优化权衡。ROLLER的设计哲学基于一个观察：大型且密集的算子往往是执行时间的主要贡献者。这导致了一个设计权衡：将优化数据复用（即最大化流水线吞吐量）作为主要优化目标，并将其他硬件相关的优化转化为对齐约束。这种权衡为大多数主流工作负载中的算子带来了快速编译和高质量的内核。对于小型算子，ROLLER进一步采用了一些自适应机制来在不同优化目标之间进行权衡，例如，使用阈值来限制冗余工作（§3.1），采用自适应的rTile收缩过程来增加并行性（§3.2）等。

未来工作。ROLLER目前依赖高级设备编译器（如nvcc）将内核代码编译为可执行文件。这有时会引入不良的性能影响，并迫使ROLLER保守地分配寄存器，因为设备编译器会隐式地为中间值（如循环变量）分配寄存器，而ROLLER无法预先检测到这种隐式分配。未来的工作之一是直接生成汇编代码（如NVIDIA GPU的PTX）以避免高级设备编译器的副作用。此外，尽管影响性能的关键硬件信息通常在硬件规格中可用，但仍有一些设备（如移动GPU）缺乏此类信息。另一个未来工作是利用一些分析技术【25, Romou: Rapidly generate high-performance tensor kernels for mobile gpus+2022+MobiCom 2022】来揭示和量化这些硬件特性。ROLLER的HAL假设硬件包含同构的计算单元和对称的内存访问，但一些设备具有NUMA架构，这使得微观性能模型难以估计rTile性能。我们把这个问题留作未来的工作。最后，稀疏内核的优化也可能违反DNN内核中同构工作负载的假设，使微观性能模型不准确。ROLLER假设更高级别的、支持稀疏性的编译器（如SparTA【34, Deep-learning model sparsity via tensorwith-sparsity-attribute+2022+OSDI’22】）将解决此问题。

7. 相关工作

张量编译器。大多数张量编译器将DNN算子视为嵌套的多层循环计算，这本质上定义了一个具有组合复杂度的巨大空间。TVM【15, TVM: An automated endto-end optimizing compiler for deep learning+2018+OSDI 18】继承了Halide【27, Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines+2013+PLDI ’13】的思想，将DNN算子描述为循环优化调度原语。AutoTVM【16, Learning to optimize tensor programs+2018+NeurIPS】通过应用机器学习方法从手动编写的代码模板中搜索最佳配置来扩展TVM。FlexTensor【35, Flextensor: An automatic schedule exploration and optimization framework for tensor computation on heterogeneous system+2020】提出无需手动模板自动探索空间。Ansor【33, Ansor: Generating high-performance tensor programs for deep learning+2020+OSDI 20】进一步推进了这种自动化，生成了更大的搜索空间并采用演化算法寻找高性能内核。Tiramisu【14, Tiramisu: A polyhedral compiler for expressing fast and portable code+2019+CGO】、AKG【32, Akg: Automatic kernel generation for neural processing units using polyhedral transformations+2021+PLDI】和Tensor Comprehensions【29, Tensor comprehensions: Framework-agnostic high-performance machine learning abstractions+2018+CoRR】等编译器应用多面体技术进行循环优化。所有这些方法都依赖于巨大的搜索空间，导致编译时间过长。ROLLER则探索了一种不同的、基于构造与硬件特征对齐的rTile的方法。

其他相关编译器与框架。Tensor Processing Primitives (TPPs)【18, Tensor processing primitives: A programming abstraction for efficiency and portability in deep learning workloads+2021+CoRR】定义了一组2D张量算子来组合复杂的高维张量算子，表达能力有限。相比之下，ROLLER不限制tile形状的维度，可应用于通用的张量表达式。OpenAI Triton【28, Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations+2019】是一个用于开发基于块的GPU内核的编程框架和编译器，但它依赖程序员决定块大小和调度，而这正是ROLLER通过同时考虑硬件特性和张量形状解决的关键问题。MLIR【5, MLIR. https://mlir.llvm.org/】 和Tensor IR【10, TensorIR. https://discuss.tvm.apache.org/ t/rfc-tensorir-a-schedulable-ir-for-tvm/ 7872】计划在其IR中支持块级（即tile）计算表示，ROLLER的rTile抽象和rProgram构造与这些计划兼容。

图级编译器与特定算子优化。像XLA【11, XLA. https://www.tensorflow.org/xla】、TVM 【15, TVM: An automated endto-end optimizing compiler for deep learning+2018+OSDI 18】和Rammer【26, Rammer: Enabling holistic deep learning compiler optimizations with rtasks+2020+OSDI 20】这样的图级DNN编译器专注于跨算子优化。ROLLER的内核生成与这些编译器兼容，并且其rTile抽象补充了Rammer中的rTask概念。一些工作专注于特定算子的优化，如CUTLASS【7, NVIDIA cutlass. https://github.com/NVIDIA/cutlass】是实现矩阵乘法的模板，还有针对卷积的分析模型 【24, Analytical characterization and design space exploration for optimization of cnns+2021+ASPLOS】和Winograd卷积优化【30, Drew: Efficient winograd cnn inference with deep reuse+2022+WWW】。ROLLER的优化方法对各种设备上的DNN算子是通用的。

A5 结论

ROLLER采用了一种非传统的深度学习编译器方法。它不依赖于昂贵的机器学习算法在巨大的搜索空间中寻找好的解决方案，而是使用一种基于递归构造的算法，利用新的rTile抽象来生成高效的内核。rTile将可选形状限制在与多种硬件特性对齐的范围内，数量远少于传统方法。构造出的程序可以通过微观性能模型进行评估，无需每次都在真实设备上运行。因此，ROLLER可以在数秒内编译出高性能内核，即使在不够成熟的加速器上也是如此。更重要的是，ROLLER为新的人工智能硬件供应商提供了一种独特且效率显著更高的方法来构建有竞争力的、特定于供应商的DNN库，从而弥合与市场领导者的生态系统差距，促进人工智能加速器的创新。