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

煤油灯科技2022-08-22 10:27

摘要

现代的张量编译器虽然取得了很多的进展,但在各种加速器(例如 GPU)上为操作员生成高效内核(深度神经网络 (DNN) 中的计算密集型子任务)通常需要花费数小时,这显着减慢了 DNN 的开发周期,并给通用内核库和自定义内核的开发带来了沉重的负担,尤其是对于新硬件供应商而言。编译过程缓慢是由于现有 DNN 编译器制定的搜索空间很大,必须使用机器学习算法才能找到好的解决方案。

在本文中,我们介绍了 ROLLER,它采用不同的基于构造的方法来生成内核。ROLLER 的核心是 rTile,这是一种新的 tile 抽象,它封装了与底层加速器的关键特性一致的张量形状,从而通过限制形状选择来实现高效执行。ROLLER 采用递归的基于 rTile 的构建算法生成基于 rTile 的程序(rProgram),其性能可以通过微性能模型进行有效评估,而无需在真实设备中进行评估。因此,ROLLER 可以在几秒钟内生成高效内核,其性能可与 GPU 等流行加速器上的最先进解决方案相媲美,同时在 IPU 等不太成熟的加速器上提供更好的内核。

1 介绍

深度神经网络越来越重要,深度学习编译器在硬件上生成高效的Kernel也越来越重要,并且取得了很多成功。但是当代的编译器在生成高效的Kernel时往往需要搜索数个小时甚至数天,因为它们都是把这些网络中的算子实现成多重循环嵌套。张量编译器通常需要对已实现的多重循环计算进行循环展开、合并、分块、缓存使用、改变并行度等调整以适应硬件的内存结构(比如CPU的三级Cache和CUDA的global memory,l2 cache, l1 cache结构)或者硬件特性(比如向量化,并行化)。这里涉及到非常大和复杂的搜索空间,所以搜索时间会很久。本文提出的Roller解决了搜索时间长的问题,它有如下几个特点:

  • 1,Roller不把DNN中的算子计算视为多层嵌套循环,而是视作数据处理管道,其中数据块(tile) 在具有并行执行单元(如GPU SM)和内存层次结构抽象的硬件上移动和处理。生成高效的Kernel的目标变成了提高流水线吞吐量的目标。

Roller将算子的计算过程建模为基于数据块(tile)的流水线,即将不同大小的数据块从多级内存结构中搬运到处理器如SM计算并逐级写回。

  • 2,为了使得基于数据块的流水线吞吐量最大化,要求每一级的数据块(Tile)shape都必须匹配(论文中叫对齐)硬件的参数设置,比如memory bank, memory transaction length, 和 minimum schedulable unit (e.g., warp size in GPUs)这些和内存带宽以及并行度相关的设置。这个约束不仅可以使得张量程序在每一级内存中都拥有很好的计算效率,这还大大降低了以前以多重循环为基础的参数搜索空间,从而解决张量编译器在编译时因为搜索Schedule耗费的大量时间。
  • 3,对齐硬件的数据处理管道的性能是高度可预测的。因为内存吞吐量可以从硬件规范或者Benchmark测试得出,这大大简化了对不同硬件进行对齐后做性能估计的难度,并不再需要基于硬件去构建复杂的代价模型来估计性能。

基于这些想法,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之上实现了Roller并开源了代码。大量的实验表明Roller可以在几秒内生产高度优化的Kernel,特别是对于大型自定义的高成本算子。这在编译时间上实现了3个数量级的改进。Roller生成的Kernel可以和最先进的张量编译器乃至硬件厂商提供的加速库相媲美,并且通常性更好(指接入新的硬件)。使用三个 rTile-based 的接口(Load, Store, Compute)描述一个程序,Roller可以轻松适应不同的加速器如AMD GPU和Graphcore IPU。

2 动机和关键观察

  • Excessive compilation time:过度的编译时间。我们自己对一组DNN操作符的经验表明,使用先进的张量编译器Ansor的单个操作符的平均编译时间为0.65小时。其中ResNet模型中一个卷积算子耗时2.17小时。一个DNN模型可能包含数百个操作符,因此很容易需要几天来编译模型。例如,为了编译一个NASNet模型,在调优41.8小时后,我们只获得了整个搜索进度的32%。我们的经验还表明,编译速度在不太成熟的设备上甚至更差,编译器在内核上花费的时间要长得多。
  • Observation and insights:我们观察到对于深度学习算子的计算有不同的视角。以矩阵乘法为例子来说明我们的观察。和将MatMul视为跨三重循环的现有编译器不同,算子的计算过程也是一个数据处理管道。我们可以从A和B Load 2个子矩阵(tile),Compute 两个子矩阵,Store 结果到C的内存中。所以计算的性能取决于 Load-Compute-Store 管道移动一个 Tile 有多快。

影响流水线中所有步骤关键性能的因素是Tile shape和一维内存空间中的布局。图1(a)说明C中一个元素的计算和内存访问的模式。假设所有矩阵存储在行优先的布局中,从B加载列会有1个跨步访问。假设这里的事务内存长度(the memory transaction length)是4,那么就有3/4的冗余数据读取。所以数据块的形状应该和内存事务长度对齐,以实现高效的内存访问。在图1(b)中,当以1x4 Tile的粒度计算B时不会有内存带宽浪费。除了内存对齐之外,数据的Tile shape还应该和硬件执行单元如并行线程数对齐以避免浪费计算周期。此外,由于Cache的存在,Tile shape也会影响数据重用机会。例如Figure1(a)每次计算1x1 tile时需要读取2mnk个数据。然而在图1(b)中只需要1.25mnk次读取,因为来自A的一次数据读取可以重复使用4次。如果沿M维度的tile 大小设置为4x4,总的reads可以减少到0.5mnk,总的数据读取效率比Figure1(a)提高了10倍。

图1

3 系统的设计

下面的图2描述了Roller的系统设计。Roller的输入是使用TE表达式,该表达式由用户生产或者从其它编译器生成(这一步可能会发生一些融合操作)。Roller从TE中提取张量形状并基于硬件规范来构建rTiles,即对齐硬件的构建块。基于rTiles,Roller提出了一种横扩纵扩递归构造算法,用于生成描述数据处理管道的高效张量化程序(rProgram)。在生成rProgram时,构建算法通过微观性能模型评估构建的rProgram的性能,从而识别出良好的rTile配置。它建立在通过硬件抽象描述的设备上,仅公开和rTiles相关的接口:Load/Save/Compute。构建的rProgram最终通过Codegen生成特定设备的最终Kernel。

图2 Roller系统概述

3.1 张量表达式和rTile

ROLLER将一个张量计算的输入作为基于lambda表达式的索引,即张量表达式。它描述了输出张量中的每个元素是如何根据输入张量中相应的元素计算出来的。例如,一个输出张量c为M ×N形状的MatMul算子可以表示为:

C = compute ( ( M, N ), lambda i , j : sum (A [ i, k ] * B[ k, j ] ) ),

其中C中索引为 (i, j )的元素通过对a第i行元素和B第j列元素的和约简计算,k为约简轴。该表达式可以覆盖DNN模型中的大部分算子,应用非常广泛,现有的DNN编译器包括TVM , Ansor和FlexTensor。

Roller引入rTile作为基本计算单元来组成张量计算。如图3所示,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,如图4所示。

图3 rTile的数据结构

图4 由rTile为MatMul表达式推断的数据tile和计算tile

rTile的一个独特属性在于它必须和给定张量表达式中的底层硬件特征和输入Tensor shape保持一致。所有这些对齐方式都由Figure3里rTile 的 shape 和 storage_padding 来控制,它们分别代表 rTile 的逻辑形式和物理布局。接下来,详细阐述对齐的详细要求:

  • Alignment with the hardware execution unit 与硬件执行单元对齐。首先,rTile的shape必须和它运行的执行单元的并行度对齐。例如,在GPU上运行 rTile 的shape 大小必须是 wrap size的倍数比如 32 来达到最大的计算效率。当在NVIDIA GPU中使用TensorCore时,rTile shape大小应该是 16x16x16 的倍数。
  • Alignment with memory transaction与内存事务的对齐 。其次,数据块(Tile)的 shape 应该和内存事务的长度保持一致,以实现最佳内存访问。具体来说,对于rTile的每个数据块我们都应该保证它的Leading dimension(如行优先Tensor中的最内层维度)是内存事务长度的倍数。如图5(a)所示,在Roller中,张量内存以缓存对齐的方式分配。因此,rTile可以避免浪费任何的内存读取,因为它的 shape 是和内存事务长度对齐的。

最大程度的利用全局内存带宽,提高全局内存加载效率是优化Kernel的基本条件,非对齐的内存会造成带宽浪费。

图 5

  • Alignment with memory bank 与存储库对齐。第三,数据块的内存布局应该和Memory Bank对齐,以避免读取冲突。例如,在图5(b)中数据块a(shape为[3, 4] )跨4个bank保存在内存中,并由形状为 [3, 1] 的块读取。将这个形状为[3, 1]的小块中的数据存储在一个bank的naive方法将导致加载冲突。rTile通过padding来避免这种低效。给定一个Leading dimension为N的数据块,由另外一个Leading dimension为n的块读取,我们延N维度做一个padding\_size大小的padding。

其中B和L分别是bank数量和bank的宽度。每一个维度的padding大小被计算出来后存到Figure3中的storage_padding字段。对于图5(b),通过padding_size为1的填充,所有的值 [3x1] 分布在不同的bank中,可以高效读取。

  • Alignment with tensor shape 与张量形状对齐。最后,rTile的形状应该与输入张量表达式的张量形状对齐。否则,rTile不能平均地划分计算,浪费计算资源或产生大量的边界检查开销。一个简单的解决方案是沿着大小为Ni的张量维i添加一个填充大小Pi,这使得Ni +Pi是rTile形状在i轴的尺寸大小的倍数。然而,一个大的填充可能会浪费计算的时间。所以,Roller将张量padding限制在ε内,并且满足以下公式:

其中Ni是第i维的张量大小。这确保了浪费的计算百分比在ε范围内。有了这个限制,我们可以枚举所有满足这个条件的有效rTile形状。

  • Deriving allrTiles 推导 allr Tiles。.鉴于上述对齐要求,对于特定的张量表达式和硬件设备,Roller 使用以下接口增量导出所有符合条件的 rTiles:

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

在给定设备指定参数后,它返回rTile shape里每个维度的下一个对齐大小。这是通过在每个维度逐渐增加尺寸大小直到满足所有对齐要求来计算的。rTile抽象允许Roller被扩展以支持新的对齐要求,这是通过GetNextAlignedAxisSize接口来实现的。

  • Calculating data reuse score rTile 计算数据重用评分rTile。一个有趣的特性是我们可以通过调整它的shape来隐式的控制内存流量,增加rTile 大小通常会以占用更多内存为代价为程序带来更多的数据重用机会,给定一个rTile T和在每一个轴上的下一个对齐大小,我们可以通过:

计算出轴i的数据重用分数S ,其中Ti_i 是通过GetNextAlignedAxisSize得到的一个对其大小替换轴i处维度大小获得的一个更大的rTile。函数Q(T)和F(T)计算以T的粒度执行计算时的内存流量和内存占用,这可以根据给定张量表达式和硬件内存规范直接推断,更大的Si_i 意味着在使用相同内存时可以节省更多的内存流量。内存重用分数在构建高效的rProgram(使用rTiles)中起着至关重要的作用。

3.2 张量项目建设

  • rTile program rTile程序。给定 rTile 和现代加速器的内存分层结构,张量计算可以自然地被看成数据流处理管道。计算从最低的内存级别加载数据块(在rTile中指定),在加速器的执行单元上对rTile进行计算,并将结果数据块写回最低的内存级别。对于每个内存级别,定义了一个特定的rTile和该内存级别的特性保持一致。 因此,Roller将张量计算描述为具有分层 rTile 配置的数据处理管道,成为rProgram。

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

图6

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

图7 Roller计算模型:(a) rTile程序;(b)矩阵乘法rTile程序;(c) rTile程序在硬件内存层次结构上的执行。

给定一个数据处理流水线,对应的rProgram的优化目标就是最大化流水线的吞吐量。这个目标可以转化为满足三个条件:1)计算和内存移动应该充分利用硬件的特性。2)吞吐量应该达到瓶颈阶段(接近峰值)。3)需要有足够的并行度来利用所有的并行执行单元。因此,Roller提出以下rProgram的构建策略:首先通过构建单核 rProgram在一个内核上纵向扩展,使得Kernel的硬件利用率饱和。然后通过复制构建的单Kernel横扩以利用硬件的并行度。

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

图8 ROLLER的单核rProgram构建算法

图8 展示了详细的构建算法。给定一个张量表达式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节的微性能模型推导出来的。

  • Scaling out an rProgram 向外扩展rProgram。鉴于大多数DNN算子的计算模式和加速器中的并行执行单元的同质性,Roller 通过将计算统一划分为大小等于最低内存层级 rTile 的 rTiles,简单地将在一个执行单元上构建的 rProgram 复制到其他单元。我们通过将所有rTiles平均分配到所有执行单元来实现这一点。注意,Roller更喜欢将reduce轴分配到同一执行单元上,因为它们可以在更高的内存层级中共享recue的结果。请注意,Roller并不假设会独占所有的计算单元,系统可以在横向扩展时显示地控制rProgram的并行度。
  • Small operator and irregular tensor shape 小算子和不规则张量形状。横向扩展算法天然有利于有足够并行度的算子。例如,分区数明显大于执行单元数的。对于小算子,算法的整体性能kennel会受到并行执行单元利用率低的影响。这里可以通过Rammer编译器的同时调度一些小Kernel来解决。然后另外一种方法是对于每个rProgram,Roller尝试沿着具有最小数据重用分数的轴收缩rTiles,来实现足够的并行度。请注意,和其它对齐规则一样,此枚举过程每次都会返回下一个对齐的Tile大小,这是一个高效的过程,和整个构建过程相比产生的成本可以忽略。

另外大算子可能包含不规则的尺寸较小的张量维度,而Roller由于对齐要求kennel无法生成足够数量的rProgram。为了解决这个问题,Roller通过轴融合pass将张量表达式转换为规范的形式。具体来说,对于所有设计的张量,如果在一个张量中存在两个相邻的轴,这些轴在所有的其它张量中既存在又相邻,或者都缺失,Roller就可以安全的合并这两个轴。如,一个输入和输出张量形状都是张量,Roller会把这三个维度fuse起来变成。除了轴融合外,Roller还尝试在张量填充机制中贪心的增加参数,直到kProgram构建完成。

3.3 rProgram的有效评估

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

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

图9 ROLLER硬件抽象的接口

  • Micro performance model 微观性能模型。借助硬件抽象层,Roller可以轻松推导出rTile(和rProgram)的性能。首先,给定一个rTile,可以从rTile的张量表达式expr和shape(图9中的MemFootprint 和 MemTraffic 接口)静态推断出产生的内存占用(包括padding)和跨不同层的内存流量。计算数据重用分数并检查rTile是否已经超出内存容量。其次,为了计算rTile的MaxComputePerf,Roller通过积极扩大Tiles shape使得TEU饱和,进行一次性分析以测量峰值计算吞吐量。此性能数据缓存在Roller中,供将来在构造算法中查询。最后,对于给定的rTile,Roller还估计MemPerf,即从内存低层加载到更高层的性能。给定rTile中对齐的内存访问,加载常规Tile的延迟可以简单地通过将总流量处于内存带宽来建模。对于所有TEU共享的内存层,我们平均分配带宽。对于较小的访问内存,Roller对每种设备类型进行一次离线分析并缓存结果。值得注意的是,微观性能模型只需要在Tile shape完全对齐的情况下准备,这是Roller的关键要求。

4 实现

  • Code generation 代码生成:给定固定的代码结构(如图6中的一个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。
  • Performance profiling 性能分析。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 仅并行分析数十个候选者。未来,我们计划实现汇编级代码生成,以缓解高级设备编译器中的隐藏问题。

5 测评

图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 倍。

图10 在V100 GPU上算子的性能表现

图11 还比较了算子编译的平均时间:

图11 每个算子的编译时间

可以看到相比于TVM和Ansor,Roller的算子编译时间在数秒内,比TVM和Ansor的搜索时间快了2个数量级。

V100 gpu上的端到端模型执行时间(以毫秒为单位)和编译时间

这个表格展示了几个经典的神经网络的性能和编译时间,可以发现Rooller相比于TVM和Ansor可以获得相当的性能,但可以将编译时间从几十个小时缩短到几百秒钟,可以大大提高模型的实际生产周期。

6 讨论与未来工作

  • Optimization space compared with loop-based compiler与基于循环的编译器相比的优化空间。rTile的抽象和数据处理管道使ROLLER构建了一个与现有的DNN编译器重叠但不同的优化空间。如前所述,这些编译器将张量编译视为嵌套循环优化。例如,Ansor只允许沿着10 或尺寸的可分割平铺大小来均匀地划分一个环轴。这使得它对于具有素数次元的张量形状的性能通常更差。相反,ROLLER专注于从数据处理管道的角度最大化硬件效率,允许更激进的优化,例如,探索不可分割但硬件对齐的平铺大小,并融合相邻轴和填充张量形状。由于我们观察到大多数DNN操作符都是内存绑定的,ROLLER与现有的DNN编译器有本质区别,它首先优化数据块吞吐量,即最大化重用分数奖励并与硬件特性保持一致,然后是并行性。这样的权衡可以带来快速编译和良好的性能,使算子具有足够的并行性。
  • Optimization trade-off优化的权衡。ROLLER的设计理念是基于一项观察:大而密集的算子往往是执行时间的主要贡献者。这就导致了一种设计权衡:优化数据重用(即最大化管道吞吐量)作为主要的优化目标,而将其他与硬件相关的优化变成对齐约束。这种权衡使得主流工作负载中的大多数操作符可以快速编译和高内核质量。对于小的算子,ROLLER进一步采用了一些自适应机制来权衡不同的优化目标,例如,当结果不足时,使用一个阈值来限制冗余工作,使用一个自适应的rTile收缩过程来增加并行度等。
  • Future work.未来的工作。ROLLER目前依赖于高级的设备编译器,例如NVCC,来编译内核代码到可执行文件。这有时会带来不良的性能影响,并迫使ROLLER保守地分配寄存器,这是因为设备编译器将隐式地为中间值分配例如,循环变量)。ROLLER无法预先检测隐式寄存器分配,因此很难估计和决定精确的寄存器使用情况。我们未来的工作之一是直接生成汇编代码(例如,NVIDIA GPU的PTX),以避免高级设备编译器的副作用。

此外,尽管影响性能的关键硬件信息,包括内存带宽、容量和事务长度,通常可以在硬件规范中找到,但仍然有一些设备(例如移动GPU)缺乏这些信息。另一个未来的工作是利用一些分析技术来揭示和量化那些硬件特性。

ROLLER的HAL假设硬件包含同构计算单元和对称的内存访问。然而,我们也观察到一些设备具有NUMA架构。这使得微性能模型很难估计rTile的性能,因为在NUMA架构下,相同的tile在不同的位置会有不同的性能。我们把这个问题留作以后的工作。

最后,对稀疏核的优化也可能会违背DNN核中均质工作量的假设,使微性能模型不准确。一些稀疏程度较大的可能与密集表现不同。ROLLER采用了更高的层次,稀疏性感知编译器(例如SparTA)将解决这个问题。

7 相关工作

大多数张量编译器将DNN运算符视为嵌套的多级循环计算,它本质上定义了一个具有组合复杂性的大空间。TVM从Halide继承了insight,并将DNN操作符描述为循环优化调度原语。后来,AutoTVM倾向于TVM应用ml方法从手动编写的代码模板中搜索最佳配置。FlexTen建议在没有手动模板的情况下自动探索空间。安索进一步推进了这种自动化。该算法考虑了代码的层次化结构,产生了更大的搜索空间,并采用进化算法来寻找性能内核。像Tiramisu、AKG和张量推导这样的编译器应用基于多面体技术来进行循环优化,它将循环转换为整数规划问题,并通过求解器找到一个良好的配置。所有这些方法都依赖于巨大的搜索空间来提供良好的内核,这导致了很长的编译/解决时间,ROLLER探索了一种不同的方法来构造与硬件特性一致的rTiles。

张量处理原语(Tensor Processing Primitives, TPPs)定义了一组2d张量运算符,在高维张量上组合复杂运算符,提供有限的表达能力。相反,ROLLER不限制tile形状的尺寸,可以应用于一般张量表达式。OpenAI Triton是一个用于开发基于块的GPU内核的编程框架和编译器。Triton依赖于程序来决定块大小和块调度,而这是ROLLER通过同时考虑硬件特性和张量形状来解决的关键问题。MLIR和Tensor IR计划在其IR中支持块级(即tile)计算表示,ROLLER的rTile抽象和rProgram构造与这些倡议是兼容的。

图级的DNN编译器,如XLA, TVM和Rammer专注于跨操作优化,例如,操作融合/协同调度。ROLLER的内核生成与这些编译器兼容,ROLLER的rTile抽象补充了Rammer中的rTask概念,因为它提供了一种构造rTask的有效方法。

最后,一些工作集中于特定于算子的优化。CUTASS是一个实现矩阵乘法的模板。针对多核CPU上的卷积算子,提出了一个解析模型来寻找最优环路级优化配置。DREW提出了一种利用数据压缩优化Winograd卷积的新方法,ROLLER的优化方法适用于不同设备上的DNN操作人员

8 结论

ROLLER对深度学习编译器采取了一种非常规的方法,与依赖昂贵的机器学习算法在一个大的搜索空间中找到一个好的解决方案不同,ROLLER使用一个基于递归构造算法生成高效的内核,该算法利用了新的rTile抽象,更少的形状与多个硬件特性对齐。构建的程序可以通过微性能模型进行评估,而无需每次在真实设备上运行。因此,即使在不太成熟的加速器中,ROLLER也可以在几秒钟内编译高性能内核。更重要的是,ROLLER为新的AI硬件供应商提供了一种独特的、明显更有效的方法,以建立有竞争力的供应商特定的DNN库,弥合生态系统与市场领导者之间的差距,从而促进AI加速器的创新。

完整附件下载:ROLLER Fast and Efficient Tensor Compilation for Deep Learning.pdf

你可能也喜欢这些文章

科技互联网,科技,资讯,动态,洞察,量子,计算,AI,人工智能,机器人,区块链,Web3,分布式,操作系统,OS,芯片,视频,深度,论文