今天来阅读一下最近 OSDI 放出的微软的 Roller 这篇论文,题目为:《Roller: Fast and Efficient Tensor Compilation
for Deep Learning》
前段时间我分享了一下 OSDI 2021 PET: Optimizing Tensor Programs with Partially Equivalent Transformations
and Automated Corrections》 这篇论文的解读。去年也分享了 OSDI 2020 《Ansor : Generating High-Performance Tensor Programs for Deep Learning》这篇论文的解读。这两篇论文的解读可以在这个地址:https://github.com/BBuf/tvm_mlir_learn/tree/main/paper_reading 或者知乎主页找到 。Ansor 的主要贡献是做到了自动寻找高效的Schedule(循环展开、合并、分块、缓存使用、改变并行度等等),不再需要开发者在TVM中基于Tensor Expression手写Schedule模板,大大增强了算子编译器(Tensor Compiler)的易用性并且对一些典型的算子和模型效果也很不错,算是AutoTVM的升级版(因为AutoTVM还需要手动指定需要搜索的Schedule模板:https://zhuanlan.zhihu.com/p/508283737)。PET则不关心算子的Schedule,而是从部分等价变换的新角度出发去增加并行度或者改善缓存从而达到加速效果,和Roller这篇论文没什么关系,其实读不读都没关系。
最近不少人问一些tvm相关的问题,我也是业余看了下所以很多时候不能很好的解答,我建立一个讨论TVM的微信群吧,有需要的读者可以加一下互相问一问。请加微信 bbuf23333 入群,备注一下tvm吧。另外业余接触编译器这一年整理的这个知识仓库已经有500+ star了,谢谢大家。希望能得到更多关注。
https://github.com/BBuf/tvm_mlir_learn
无论是Ansor,AutoTVM还是PET(一部分代码生成也是基于TVM AutoTVM/Ansor的)它们都面临了同样一个问题,那就是在对算子的Schedule进行搜索时需要耗费大量的时间,在特定硬件上对一个常见的视觉模型进行自动调优和生成代码kennel需要数小时。这严重阻碍了AI编译器应用于模型部署。基于这个痛点,Roller横空出世。
ROLLER:一个用于深度学习的快速高效的张量编译器。作者来自微软亚洲研究院以及多伦多大学等多所高校。
现代的张量编译器虽然取得了很多的进展,但通常这些编译器都需要小时计的时间去搜索和生成高效的Kernel,这是因为现有张量编译器通常指定的搜索空间很大。为了解决编译时间长的问题,本文提出了Roller,它的核心是rTile,这是一种新的Tile抽象,它封装了和底层加速器的关键特性一致的张量shape,从而通过限制shape的选择实现高效的运行。Roller采用了基于rTile的递归构建算法来生成目标程序(rProgram)。最终,Roller可以在几秒内就生产高效的Kernel,性能可以媲美目前主流加速器上的其它张量编译器,并且为IPU等新的加速器生产更好的Kernel。
还不能看出什么,继续往下看吧。这里说的tile就是对输入进行分块以适应硬件的内存结构,我在之前的文章有详细讲到,不了解的同学可以先看一眼tile这部分的科普:https://zhuanlan.zhihu.com/p/508283737 。
深度神经网络越来越重要,深度学习编译器在硬件上生成高效的Kernel也越来越重要,并且取得了很多成功。但是当代的编译器在生成高效的Kernel时往往需要搜索数个小时甚至数天,因为它们都是把这些网络中的算子实现成多重循环嵌套。张量编译器通常需要对已实现的多重循环计算进行循环展开、合并、分块、缓存使用、改变并行度等调整以适应硬件的内存结构(比如CPU的三级Cache和CUDA的global memory,l2 cache, l1 cache结构)或者硬件特性(比如向量化,并行化)。这里涉及到非常大和复杂的搜索空间,所以搜索时间会很久。这篇文章提出的Roller解决了搜索时间长的问题,它有如下几个特点。
Roller将算子的计算过程建模为基于数据块(tile)的流水线,即将不同大小的数据块从多级内存结构中搬运到处理器如SM计算并逐级写回。
基于这些想法,Roller提出了rTile,这是一种新的抽象,它封装了和硬件加速器的关键特征和输入张量shape一致的数据块(Tile)shape(后面会详细看)。然后将数据处理管道描述为基于rTile的程序(rProgram),由Load, Store, Compute 三个接口组成,作用于rTile。为了构建高效的rProgram,Roller遵循了一个scale-up-then-scale-out的方法。它首先执行Scale-up的过程,该过程采用基于rTile的递归构造方法(Figure8)逐渐增加rTile shape大小,来构造一个饱和加速器单个执行单元(如SM)的rProgram。然后执行Scale-out的过程,由于深度学习的计算模式和加速器的并行执行单元的同质性,它只是将生成的rProgram复制到其它并行执行单元。这里的scale-up-then-scale-out可以叫做纵扩和横扩。
Roller可以在没有显著开销的情况下评估不同rTiles的性能。每种算子可以简单的测试一下峰值和带宽。由于对齐了硬件结构,其它关键的性能因素比如rTile的内存压力可以从硬件规则分析得到。这样就得到了一个高效的微评测模型,避免了其它编译器所需的对每个配置进行昂贵的在线分析,从而显著加速了编译过程。此外,由于严格的对齐要求,递归构造过程可以快速生产一些想要的rTiles和rPrograms。综合一下,Roller可以在几秒内生成高效的Kernel。
作者团队在TVM和Rammer(Rammer可以看:https://www.msra.cn/zh-cn/news/features/osdi-2020-rammer)之上实现了Roller并开源了代码。大量的实验表明Roller可以在几秒内生产高度优化的Kernel,特别是对于大型自定义的高成本算子。这在编译时间上实现了3个数量级的改进。Roller生成的Kernel可以和最先进的张量编译器乃至硬件厂商提供的加速库相媲美,并且通常性更好(指接入新的硬件)。使用三个 rTile-based 的接口(Load, Store, Compute)描述一个程序,Roller可以轻松适应不同的加速器如AMD GPU和Graphcore IPU。
影响流水线中所有步骤关键性能的因素是Tile shape和一维内存空间中的布局。Figure1(a)说明C中一个元素的计算和内存访问的模式。假设所有矩阵存储在行优先的布局中,从B加载列会有1个跨步访问。假设这里的事务内存长度(the memory transaction length)是4,那么就有3/4的冗余数据读取。所以数据块的形状应该和内存事务长度对齐,以实现高效的内存访问。在Figure1(b)中,当以1x4 Tile的粒度计算B时不会有内存带宽浪费。除了内存对齐之外,数据的Tile shape还应该和硬件执行单元如并行线程数对齐以避免浪费计算周期。此外,由于Cache的存在,Tile shape也会影响数据重用机会。例如Figure1(a)每次计算1x1 tile时需要读取2mnk个数据。然而在Figure1(b)中只需要1.25mnk次读取,因为来自A的一次数据读取可以重复使用4次。如果沿M维度的tile 大小设置为4x4,总的reads可以减少到0.5mnk,总的数据读取效率比Figure1(a)提高了10倍。
下面的Figure2描述了Roller的系统设计。Roller的输入是使用TE表达式。该表达式由用户生产或者从其它编译器生成(这一步可能会发生一些融合操作)。Roller从TE中提取张量形状并基于硬件规范来构建rTiles,即对齐硬件的构建块。基于rTiles,Roller提出了一种横扩纵扩递归构造算法,用于生成描述数据处理管道的高效张量化程序(rProgram)。在生成rProgram时,构建算法通过微观性能模型评估构建的rProgram的性能,从而识别出良好的rTile配置。它建立在通过硬件抽象描述的设备上,仅公开和rTiles相关的接口:Load/Save/Compute。构建的rProgram最终通过Codegen生成特定设备的最终Kernel。
Roller将TVM中引入的Tensor Expression引入作为编译器的输入,Tensor Experssion这里不讲了,如果不了解可以看一下TVM里面chen tianqi写的文档。https://tvm.apache.org/docs/tutorial/tensor_expr_get_started.html
Roller引入rTile作为基本计算单元来组成张量计算。如Figure3所示,rTile封装了沿给定张量表达式的expr的每个循环轴定义的多维tile shape。给定shape和expr,rTile可以静态推断所涉及的输入和输出数据块。例如,沿轴i, j, k的tile shape表示上述Matmul表达式的rTile,其中每个rTile加载来自A的4x2个数据以及来自B的2x4个数据,进行总共4x2x4 次 mul-add计算,并将4x4的数据tile写回到C,如Figure4所示。
rTile的一个独特属性在于它必须和给定张量表达式中的底层硬件特征和输入Tensor shape保持一致。所有这些对齐方式都由Figure3里rTile 的 shape 和 storage_padding 来控制,它们分别代表 rTile 的逻辑形式和物理布局。接下来,详细阐述对齐的详细要求:
最大程度的利用全局内存带宽,提高全局内存加载效率是优化Kernel的基本条件,非对齐的内存会造成带宽浪费,可参考:https://face2ai.com/CUDA-F-4-3-%E5%86%85%E5%AD%98%E8%AE%BF%E9%97%AE%E6%A8%A1%E5%BC%8F/
其中B和L分别是bank数量和bank的宽度。每一个维度的padding大小被计算出来后存到Figure3中的storage_padding字段。对于Figur5(b),通过padding_size为1的填充,所有的值 [3x1] 分布在不同的bank中,可以高效读取。
GPU Shared Memory bank conflict: https://blog.csdn.net/Bruce_0712/article/details/65447608
Alignment with tensor shape 最后,rTile的shape应该和输入张量表达式的张量shape对齐。否则,计算不能被rTile均匀且分,浪费计算资源或者产生大量的边界检查开销。一个简单的解决方案是沿着Tensor的维度 i i i(大小为 N i N_i Ni)进行padding,padding的大小为 P i P_i Pi,使得 N i + P i N_i+P_i Ni+Pi时rTile shape在维度i大小的倍数。但是较大的padding kennel会浪费计算,所以Roller将张量padding限制在 ε \varepsilon ε内,并且需要满足以下公式: S i − N i m o d S i N i < = ε \frac{S_i-N_i \mod S_i }{N_i}<= \varepsilon NiSi−NimodSi<=ε。这确保了计算的浪费百分比以 ε 为上界。 有了这个限制,我们可以枚举所有满足这个条件的有效 rTile 形状。
Deriving allrTiles. 鉴于上述对齐要求,对于特定的张量表达式和硬件设备,Roller 使用以下接口增量导出所有符合条件的 rTiles:
vector<int> GetNextAlignedAxisSize(rTile T, Dev d),
在给定设备指定参数后,它返回rTile shape里每个维度的下一个对齐大小。这是通过在每个维度逐渐增加尺寸大小直到满足所有对齐要求来计算的。rTile抽象允许Roller被扩展以支持新的对齐要求,这是通过GetNextAlignedAxisSize
接口来实现的。
GetNextAlignedAxisSize
得到的下一个对齐大小替换轴 i i i处维度大小获得的一个更大的rTile。函数Q(T)和F(T)计算以T的粒度执行计算时的内存流量和内存占用,这可以根据给定张量表达式和硬件内存规范直接推断(0x3.3节内容)。更大的 S i S_i Si意味着在使用相同的内存时可以节省更多的内存流量。内存重用分数在构建高效的 rProgram(使用 rTiles)中起着至关重要的作用。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横扩以利用硬件的并行度。
Figure8展示了详细的构建算法。给定一个张量表达式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′)是基于dev和0x3.3节的微性能模型推导出来的。
另外大算子可能包含不规则的尺寸较小的张量维度,而Roller由于对齐要求kennel无法生成足够数量的rProgram。为了解决这个问题,Roller通过轴融合pass将张量表达式转换为规范的形式。具体来说,对于所有设计的张量,如果在一个张量中存在两个相邻的轴,这些轴在所有的其它张量中既存在又相邻,或者都缺失,Roller就可以安全的合并这两个轴。如,一个输入和输出张量形状都是[17, 11, 3]的张量,Roller会把这三个维度fuse起来变成 [ 561 ] ( 17 × 11 × 3 ) [561](17\times 11\times 3) [561](17×11×3)。除了轴融合外,Roller还尝试在张量填充机制中贪心的增加参数 ε \varepsilon ε,直到kProgram构建完成。
在构建算法中,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获取。
还有一些NIVIDIA GPU/AMD ROCm/Grphcore IPUs具体硬件上的一些实现细节,感兴趣的可以自己看下论文。
这里主要看一下在cuda上的结果。
Figure 10 绘制了我们基准测试中 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 倍。
下面的Figure11还比较了算子编译的平均时间:
可以看到相比于TVM和Ansor,Roller的算子编译时间在数秒内,比TVM和Ansor的搜索时间快了2个数量级。
这里的Table展示了几个经典的神经网络的性能和编译时间,可以发现Rooller相比于TVM和Ansor可以获得相当的性能,但可以将编译时间从几十个小时缩短到几百秒钟,可以大大提高模型的实际生产周期。
为了解决编译时间长的问题,这篇论文提出了Roller,它的核心是rTile,这是一种新的tile抽象,它封装了和底层加速器的关键特性一致的张量shape,从而通过限制shape的选择实现高效的运行。Roller采用了基于rTile的递归构建算法来生成目标程序(rProgram)。最终,Roller可以在几秒内就生产高效的Kernel,性能可以超越目前主流加速器上的其它张量编译器,并且为IPU等新的加速器生产更好的Kernel。