《Roller: Fast and Efficient Tensor Compilation for Deep Learning》

用于深度学习快速高效的张量编译器

作者微软亚洲研究院以及多伦多大学等多所高校

摘要

当前编译为了产生高效的kernel时，搜索空间大，通常使用机器学习的方法找到最优的方案，导致编译器时间长，通常需要几天甚至几周的时间。

ROLLER，使用基于构造的方法产生kernel，速度快，仅需要数秒。

提出rTile, 一种新的tile抽象，封装了tensor shape, 和底层加速器关键特征一致。

使用基于rTile的算法递归构造rProgram

tile就是对输入进行分块以适应硬件的内存结构

在这里插入图片描述

1. 前言

$[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-n2KfGErQ-1677663189104)(C:\Users\EDY\AppData\Roaming\Typora\typora-user-images\image-20230228095807393.png)]$

DNN 编译器(cuDNN cuBLAS)是为了产生更高效的kernel, 减少人工kernel库的开发。编译优化通常在循环计算中使用分块/融合/重排等，但是编译时间长。

本文基于以下三个观点提出ROLLER

1 相比多级嵌套，ROLLER将DNN算子的计算作为数据处理管道，数据块在抽象的硬件（具有并行执行单元和多层内存结构）中移动和处理

2 为了让加速器更加高效，数据块的shape应该和硬件特征（memory bank,memory transaction length, and minimum schedulable unit ）保持一致

3 pipline对齐后的的性能是可预测的

提出 rTile，rProgram posed by three interfaces: Load, Store, and Compute

2. 动机

超长的编译时间

Ansor单算子平均编译时长0.65小时，其中ResNet中一个卷积算子编译时长2.17小时。一个DNN网络可能有上百个算子，需要几天的时间编译。比如 NASNet 模型，编译时间41.8小时，获得了38%的提升。

观察和发现

在这里插入图片描述

3. 系统设计

Roller的输入是使用TE表达式。该表达式由用户生产或者从其它编译器生成（这一步可能会发生一些融合操作）。

Roller从TE中提取张量形状并基于硬件规范来构建rTiles，即对齐硬件的构建块。基于rTiles，Roller提出了一种横扩纵扩递归构造算法，用于生成描述数据处理管道的高效张量化程序（rProgram）。

在生成rProgram时，构建算法通过微观性能模型评估构建的rProgram的性能，从而识别出良好的rTile配置。它建立在通过硬件抽象描述的设备上，仅公开和rTiles相关的接口：Load/Save/Compute。构建的rProgram最终通过Codegen生成特定设备的最终Kernel。

在这里插入图片描述

3.1 张量表达式和 rTile

Roller将TVM中引入的Tensor Expression引入作为编译器的输入

在这里插入图片描述

rTile封装了给定张量表达式的expr的每个循环轴定义的多维tile shape。给定shape和expr，rTile可以静态推断所涉及的输入和输出数据块。

在这里插入图片描述

与hardware execution unit对齐

rTile的shape必须和它运行的执行单元的并行度对齐。例如，在GPU上运行 rTile 的shape 大小必须是 wrap size的倍数，比如 32 来达到最大的计算效率。当在NVIDIA GPU中使用TensorCore时，rTile shape大小应该是 16x16x16 的倍数。

在这里插入图片描述

与memory transaction对齐

对于rTile的每个数据块我们都应该保证它的Leading dimension（如行优先Tensor中的最内层维度）是内存事务长度的倍数。如Figure5(a)所示，在Roller中，张量内存以缓存对齐的方式分配。因此，rTile可以避免浪费任何的内存读取，因为它的 shape 是和内存事务长度对齐的。
在这里插入图片描述

与memory bank 对齐

避免读取冲突。例如，在数据块a（shape为[3, 4] ）跨4个bank保存在内存中，并由形状为 [3, 1] 的块读取。将这个形状为[3, 1]的小块中的数据存储在一个bank的naive方法将导致加载冲突。rTile通过padding来避免这种低效。给定一个Leading dimension为N的数据块，由另外一个Leading dimension为n的块读取，我们延N维度做一个padding_size大小的padding。

在这里插入图片描述

与tensor shape对齐

rTile的shape应该和输入张量表达式的张量shape对齐。否则，计算不能被rTile均匀且分，浪费计算资源或者产生大量的边界检查开销。

$[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-bf8ipqRV-1677663189108)(C:\Users\EDY\AppData\Roaming\Typora\typora-user-images\image-20230227212211830.png)]$

Deriving all rTiles

vector<int> GetNextAlignedAxisSize(rTile T, Dev d)

计算数据重用分数

更大的Si 表示在使用相同的内存时可以节省更多的内存流量。内存重用分数在构建高效的 rProgram（使用 rTiles）中起着至关重要的作用。
在这里插入图片描述

3.2 张量程序构造

rTileProgram

给定 rTile 和现代加速器的内存分层结构，张量计算可以自然地被看成数据流处理管道。

Figure6 展示了具有三个存储层（L0，L1，L2）的设备上的rProgram。rProgram由每个内存层次的 rTile 和 rTile 指令(Load, Store, Compute) 来描述。

在这里插入图片描述

Figure7(a)展示了Figure7(b)对应的MatMul程序。Figure7©说明了rProgram如何映射到设备的每个内存层次。具体来说，每次它从内存L2中加载一个A的4x4小块和B的4x8小块到L1中。然后从L1中加载一个A的2x1和B的1x2小块到L0(寄存器)中。每次计算完成后，结果的2x2小块会直接从L0写回到L2。

给定一个数据处理流水线，对应的rProgram的优化目标就是最大化流水线的吞吐量。这个目标可以转化为满足三个条件：

1）计算和内存移动应该充分利用硬件的特性

2）吞吐量应该达到瓶颈（接近峰值）

3）需要有足够的并行度来利用所有的并行执行单元

因此，Roller提出以下rProgram的构建策略：首先通过构建单核 rProgram在一个内核上纵向扩展，使得Kernel的硬件利用率饱和。然后通过复制构建的单Kernel横扩以利用硬件的并行度。

Scaling up an rProgram

由于rTile的对齐属性确保了硬件的效率，Roller可以只专注于通过构建正确的rTile shape来最大化每个内存层次的吞吐量。通过利用0x3.1节中定义的数据重用分数，单核rProgram构建算法从初始化rTile开始，并逐渐将其扩大到rTile中收益最大的轴（也即具有最大重用分数的）。注意，构造算法不需要精确的数据重用分数，它只是选择最大的一个来最大化吞吐量。在此过程中，内存的性能会提高直到达到计算峰值或者最大的内存容量。上述过程从上到下对每个内存层次进行重复，直到构建出所需的rProgram。请注意，如果某些张量表达式的数据重用分数保持不变，比如elemetwise算子，Roller将只为顶层构建rTiles并从底层内存内存加载它们。

在这里插入图片描述

给定一个张量表达式expr和目标设备dev，该算法在顶层内存构造一个初始化的rTile T并递归的放大T（对应第4行的EnlargeTile）。

每一步，它都会枚举下一个更大的rTile T‘，最大程度的提高数据重用得分（对应第10行的GetNextRTileShapes）。

如果T’达到内存容量（第13行）或者数据块加载的吞吐量MemRef(T’)超过了峰值计算吞吐量 MaxComputePer f(T’)（第17行），

算法记录当前的rTile并在下一个内存级别继续EnlargeTile。否则，它会在当前内存层级继续扩大T’（第20行）。

构建在最低的内存层级完成（第6行），产生一个结果并重复运行直到产生K个rPrograms（来包容编译器的隐藏因素影响）

注意，MemPer f(T′)和MaxComputePer f(T′)是基于3.3节的微评测模型推导出来的。

Scaling out an rProgram

考虑大多数DNN算子计算模式和加速器中并行执行单元的的同质性, Roller 简单地将在一个执行单元上构建的 rProgram 复制到其他单元。

Roller 通过将计算统一划分为大小等于最低内存层级 rTile 的 rTiles。

通过将所有rTiles平均分配到所有执行单元来实现这一点。注意，Roller更喜欢将reduce轴分配到同一执行单元上，因为它们可以在更高的内存层级中共享recue的结果。

采用基于rTile的递归构造方法（Figure8）逐渐增加rTile shape大小，来构造一个饱和加速器单个执行单元（如SM）的rProgram

Small operator and irregular tensor shape

对于小算子，算法的整体性能kennel会受到并行执行单元利用率低的影响。这里可以通过Rammer编译器的同时调度一些小Kernel来解决。然后另外一种方法是对于每个rProgram，Roller尝试沿着具有最小数据重用分数的轴收缩rTiles，来实现足够的并行度。

大算子可能包含不规则的尺寸较小的张量维度，而Roller由于对齐要求kennel无法生成足够数量的rProgram。为了解决这个问题，Roller通过轴融合pass将张量表达式转换为规范的形式。具体来说，对于所有设计的张量，如果在一个张量中存在两个相邻的轴，这些轴在所有的其它张量中既存在又相邻，或者都缺失，Roller就可以安全的合并这两个轴。如，一个输入和输出张量形状都是[17, 11, 3]的张量，Roller会把这三个维度fuse起来变成[ 561 ] ( 17 × 11 × 3 ) 。

3.3 Efficient Evaluation of an rProgram

在构建算法中，Roller需要评估rProgram的性能。Roller无需评估真实硬件设备中端到端的rProgram，只需要评估 rTile 的性能，如Figure8中的MemPerf和MaxComputePerf。

Roller针对硬件抽象层（HAL）中描述的设备构建了一个微评测模型。HAL将加速器建模为具有分层内存的多个并行执行单元，HAL公开了三个基于rTile的接口：Load，Save，Compute。执行单元被抽象为rTile Execution Unit（TEU），它通过Compute接口对数据块进行计算。可以将多个TEUs组织为一个组，它们可以协同加载和存储Tiles。HAL将不同的内存层（如寄存器，共享内存，DRAM）视为一种统一类型，暴露了影响Tile性能的硬件规范。硬件规范包括内存容量，事务长度，缓存行大小，Memory Banks数量，可以通过Figure9的getDeviceSpec获取。

在这里插入图片描述

微评测模型

借助硬件抽象层，Roller可以轻松推导出rTile（和rProgram）的性能。

首先，给定一个rTile，可以从rTile的张量表达式expr和shape静态推断出产生的内存占用（包括padding，Figure9中的MemFootprint）和跨不同层的内存流量（Figure9中的MemTraffic 接口）。计算数据重用分数并检查rTile是否已经超出内存容量。

其次，为了计算rTile的MaxComputePerf，Roller通过积极扩大Tiles shape使得TEU饱和，进行一次性分析以测量峰值计算吞吐量。此性能数据缓存在Roller中，供将来在构造算法中查询。

最后，对于给定的rTile，Roller还估计MemPerf，即从内存低层加载到更高层的性能。给定rTile中对齐的内存访问，加载常规Tile的延迟可以简单地通过将总流量除以内存带宽来建模。对于所有TEU共享的内存层，我们平均分配带宽。对于较小的访问内存，Roller对每种设备类型进行一次离线分析并缓存结果。

值得注意的是，微观性能模型只需要在Tile shape完全对齐的情况下准备，这是Roller的关键要求。

4 实现

实现基于TVM和Rammer,

ROLLER实现机制包含：expression optimization, construction algorithm, micro-performance model

输入是一个ONNX图或TensorFlow图。
首先ROLLER利用Rammer进行图优化；
然后ROLLER为每个(融合的)运算符派生TVM张量表达式，并通过构造算法生成相应的rProgram，进行kernel生成；
最后，生成的内核被注入Rammer的运行时并生成端到端模型代码

代码生成（Code generation）

给定固定的代码结构（如Figure6中的一个rProgram），Roller基于预定义的模板生成代码（TVM 内置调度原语）。在每个内存层级加载和存储数据块由 TVM 的 cache_read 和 cache_write 原语实现。rTile 上的分区是通过 split 和 fuse 完成的。一些rTile计算的原语是通过TVM内置API完成的。基于模板，给定的rProgram可以直接生成cuda代码。

张量填充（Tensor padding）

Roller依靠张量padding将rTiles和张量shape对齐。在实践中，最底层内存（例如 DRAM）中的大多数张量是由外部程序（例如 DNN 框架）分配的，因此我们只需在上层内存（例如共享内存）中应用padding。 Roller的张量padding目前需要输入张量表达式来指定它是否允许填充，以及默认的填充值（例如，0 表示 MatMul 运算符）。对于 Memory Bank 对齐的storage padding，我们利用 TVM 的 storage_align 原语添加padding。

性能分析器

Roller实现了两个性能分析器。

微性能分析器：通过micro-benchmark生成内存带宽，计算吞吐量等硬件指标。这是针对每种设备类型和张量表达式的一次离线分析。

内核分析器：描述了top K个kPrograms中最快的kernel.如果k大于1则用于每一个编译结果。在实际应用中，特定内核代码的性能也会受到设备编译器和硬件相关隐藏因素的轻微影响，Roller 几乎无法控制。这些因素包括不同指令类型的指令密度、寄存器分配行为、设备编译器优化、warp 调度开销等。特别是在 NVIDIA GPU 上，Roller 依靠 nvcc 将生成的 CUDA 代码编译成机器代码。但是，nvcc 的专有优化可能会对程序执行行为产生不利影响。因此，Roller 利用内核分析器快速评估性能最佳的 rProgram 并选择最佳的。较大的 K 通常可以提高kernel质量。在评估前 10、20 和 50 个结果后，我们的经验表明，前 10 名可以获得大多数情况下的最佳结果。请注意，Roller 的内核分析器不同于以前编译器中由机器学习算法驱动的评估过程。基于 ML 的方法通常需要数百甚至数千个顺序评估步骤，而 ROLLER 仅并行分析数十个候选者。未来，我们计划实现汇编级代码生成，以缓解高级设备编译器中的隐藏问题。

ROLLER on NVIDIA CUDA GPUs

The transaction length at the global memory layer is set to 32 Bytes

the bank number is 32 Bytes

the bank length is 4 Bytes

在这里插入图片描述

ROLLER on AMD ROCm GPUs

ROLLER on Graphcore IPUs

5 评测

和TVM Ansor相比，编译时间是数量级的减少，最耗时的算子需要43s, 同时其他算子耗时13s.
ROLLER 获得了当前最好的性能，甚至50%的情况下优于其他编译加速器
针对小size和不规则的shape, 效果不是很明显，是由于很难和硬件参数对齐。但本身这些算子的耗时也较少（小于1ms）
我们进行了广泛的评估，覆盖不同加速器上的不同算子（共119个算子）

Experimental setup

Benchmarks.

在这里插入图片描述

5.1 NVIDIA GPUs评估

算子性能

cuda上的结果，119 个算子的平均kernel性能，按算子类型和 ID 排序。我们将大型算子（例如，kernel时间大于 5ms）绘制在 y 轴为对数尺度的顶部子图中，而底部 4 个子图是其它中小型算子。

首先，与 CUDA 库 (CudaLib) 相比，Roller 可以为 81.5% 占比的算子获得可比的性能（即在 10% 以内），并且对于 59.7% 的算子来说甚至更快。

我们观察到，Roller 表现较差的大多数算子是具有 3×3 或更大滤波器的卷积算子，它们通常在 cuDNN 中使用更有效的数值算法（例如，Winograd [23]）来实现，并且难以用张量表示表达。这就是在这些情况下 Ansor 和 TVM 也比 CudaLib 慢的原因。

其次，与 TVM 和 Ansor 相比，Roller 也可以分别为 72.3% 和 80.7% 占比的算子获得可比的性能。其余的 27.7% 和 19.3% 主要是小算子或张量形状不规则，难以与硬件对齐。然而，这些算子的kernel执行时间通常相对较短，例如平均只有 1.65 毫秒和 1.16 毫秒。在所有算子的 54.6% 和 65.5% 占比中，Roller 甚至可以分别比 TVM 和 Ansor 生成更快的kernel。我们观察到这些算子中的大多数都是大型且耗时的。正如上面的子图所示，当算子大于 5 毫秒（最高 343 毫秒）时，Roller 可以为这些算子中的大多数实现更好的性能，例如，与 TVM 和 Ansor 相比，平均速度提高了 1.85 倍和 1.27 倍。
在这里插入图片描述