Closed KuangjuX closed 11 months ago
张量编译器用来优化算子并生成可以在特定加速器执行的代码。然而优化算子涉及到嵌套循环拆分/融合/重排序,实际上是一个组合优化问题,搜索空间巨大。目前先进的编译器像 TVM,Ansor,Flextensor 等通过机器学习算法搜索解决方案,然而这非常耗费时间。在本篇文章中提出了 Roller。
因此 Roller 提出了 rTile 抽象用于管理 data tile shapes 和加速器关键特性和输入的张量大小对齐。一系列流水线处理过程可以被描述为 rTile-based program。有 Load,Stroe,Compute 三个接口。它使用先扩展然后后扩展的方法。
C_{m,n} = A_{m,k} * B_{k,n}
$,在存在的编译器中将这个视为三重循环。但是本文可以将其视为:1. Load 子矩阵(i.e. a tile)从 A 矩阵和 B 矩阵中。2. Compute 两个矩阵。3. 存储 C 的 tile 到内存中。因此计算结果的速度依赖于数据 tiles 可以在 Load-Compute-Store 流水线中有多快。影响性能的关键因素在于 tile 的 shape 以及内存布局。图 1(a)(b)(c) 分别展示了不同的 tile 针对内存布局带来的性能的影响。
Roller 从张量表达式中提取张量 shape 并基于硬件特性构建 rTiles,随后基于递归算法生成有效率的张量程序(rProgram),构建后通过 micro-performance 模型评估性能。rTile 抽象层仅有 Load, Compute 和 Store 三个接口。
Raller 将 tensor 表达式计算作为一个基于索引的匿名表达式:$C = compute((M, N), lambda i,j: sum(A[i,k] * B[k,j]))
$,可以覆盖大部分 DNN 模型中的算子。
给定 shape
和 expr
,rTile 可以静态推断出输入和输出的 data tiles。rTile 的独特特性是他必须根据给定的张量表达式对齐潜在的硬件特性和张量形状。这些对齐由图 3 中的 rTile shape
和 storage_padding
来控制,分别标明 rTile 的逻辑形式和物理布局。
vector<int> GetNextAlignedAxisSize(rTile T, Dev d)
的接口通过给定的 rTile 以及设备用来生成 rTile 下一个对齐的 shape 的每一个 axis 大小,这是通过不断增长 dim 来查看是否满足对齐要求。rTile program 给定一个 rTile 以及一个现代加速器的内存层级,一个 tensor 的计算可以自然地被定义为一个数据流处理流水线。
Figure 6 和 Figure 7 是一个示例,首先从 L2 中加载 [4,4] 和 [4,8] 的 data tile 到 L1,然后从 L1 加载 [2,1] 和 [1,2] 的 data tile 加载到 L0,L0 计算后直接将 [2,2] 的 data tile 存储到 L2 中。
给定一个 data process pipeline 中,优化目标是最大化 pipeline 的吞吐量。目标可以被翻译成三个状态:
Scaling up an rProgram. 由于 rTile 保证了内存对齐特性,因此 Roller 可以通过正确的 rTile 来最大化内存重用率。单核心 rProgram 构建逐渐扩大 rTile 直到 most-effective cost 的轴。
Figure 8 展示了构建 rTile 的算法,输入张量表达式和设备,通过递归调用 EnlargeTile
来获得 rTile,每次将会获得 Top K 个,防止设备有些隐藏特性不能使用。
Scaling out an rProgram. Small operator and irregular tensor shape.
Roller 也引入了一个 efficient evaluation 用来评估性能,在 Figure 8 中提到了。这些 evaluation 可以通过一次 offline profile 并将其缓存。
Roller 的实现基于 Rammer 和 TVM 两个开源项目。
storage_align
原语来添加 padding。
paper: https://www.usenix.org/system/files/osdi22-zhu.pdf