Roller
Zhu, Hongyu, et al. "ROLLER: Fast and efficient tensor compilation for deep learning." 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22). 2022.
https://www.usenix.org/conference/osdi22/presentation/zhu
1. Introduction
- Insight
- 视 DNN 算子的计算过程为数据处理流水线（data processing pipeline）。目标：提升流水线吞吐量
- 数据 tile 应该与硬件特征对齐。最大化流水线吞吐量，即构造对齐 tile 以使得计算单元饱和
- 对齐流水线的性能高度可预测。对齐配置下，简化性能评估，无需复杂的 cost 模型与高昂的硬件评估。
- rTile：与硬件特征和输入形状对齐的数据 tile 形状的封装
- Load，Compute，Store
-
rProgram：数据处理流水线可表述为基于 rTile 的程序，即 rProgram
- Scale-up 过程：递归的基于 rTile 的构造算法，逐步增加 rTile 的大小来构造使得计算单元（SM）饱和的 rProgram
- Scale-out 过程：基于深度学习的就算模式和家属器中的并行执行单元的同质性，复制 rProgram 到其他执行单元。
-
评估：micro-performance model
- 巅峰计算吞吐量可测
- rTile 的内存压力可由硬件特征中分析得出
2. Motivation and Key Observations
- 高昂编译时间
- 观察和见解
3. System Design
-
算子描述为张量表达式（tensor expression）
- 从图级 DNN 编译器 or 用户 生成
- 可能将多个算子融合
-
利用硬件特征构造 rTIle
- 利用 scale-up-then-scale-out 递归构造 rProgram
- 通过 micor-performance model 评估 rProgram 的性能
- 建立在硬件抽象层（hardware abstraction layer）之上
- 硬件抽象层只暴露 Load，Compute，Store 接口
3.1 张量表达式与 rTile
-
张量表达式：基于索引的 lambda 表达式
-
rTile：RollingTile，组成张量计算的基本单元
class rTile{ TensorExpor expr; TileShape shape; TIleShape storage_padding; vector<TileShape> GetInputDataTiles(); vector<TileShape> GetOutputDataTile(); }- shape 定义了对于给定张量表达式 expr 的每个循环轴的 tile 形状
- 根据 expr 和 shape，可静态推断涉及的输入输出数据 tile
- 沿着 MatMul 的 i,j,k 三轴的 shape [4,4,2]：表示从 A 加载 \(4\times2\) 大小，从 B 加载 \(2\times4\) 大小，进行 \(4\times4\times2\) 的乘加运算，将 \(4\times4\) 的 tile 存储到 C。
-
rTile 特性：对齐（由 shape 和 storage_padding 控制）
- 与硬件执行单元对齐
- 若在 thread warp 上运行，rTile 应该是 32 的倍数
- 若在 TensorCore 上运行，rTile 应该是 \(16\times16\times16\) 的倍数
- 与内存事务对齐
- rTile 的每个数据 tile 应保证 row-major 张量的最内层的维度（leading dimension）是内存事务长度的倍数
- 和内存 bank 对齐
- data tile 的存储布局的步长应该与存储 bank 对齐，以避免读取冲突
- 填充一个 padding size 以避免这种冲突
- 与张量形状对齐
- 若计算无法按 rTile 平均分配，将将造成计算资源的浪费，且产生高昂的边界检测开销
- 在范围 \(\epsilon\) 内填充 padding size，
- 与硬件执行单元对齐
-
给出所有的 rTile
-
计算数据重用分数（data reuse score）
-
通过调整 shape 隐式调整 memory traffic
-
数据重用分数 $$ S_i=\frac{Q(T)-Q(T_i')}{F(T_i')-F(T)} $$
- \(T_i'\) 为通过将轴 i 替换为
GetNextAlignedAxisSize中给定的值，从而新放大的 rTile， - \(Q(T)\) memory traffic
- \(F(T)\) memory footprint
- \(T_i'\) 为通过将轴 i 替换为
-
\(S_i\) 越大，cost-efficiency 越高，在相同的内存使用上，节省更多的内存流量。
-
3.2 张量程序构造
rTile 程序（rProgram）
-
流数据处理流水线：自底向顶 load 数据，然后在执行单元进行 rTile 计算结果存回最底层 memory。
-
针对每层存储层的特性，定义一个对齐的特定 rTile。
-
例子：三层存储层
-
对于给定的数据处理流水线，rProgram 的优化目标是：最大化流水线的吞吐率
- 计算和内存传输充分利用硬件功能
- 吞吐量应该使瓶颈阶段饱和
- 足够的并行性来利用所有的并行执行单元
-
构造策略：
- scale-up：通过构造一个单核 rProgram 来使得核的硬件利用率饱和
- scale-out：通过复制构造好的单核 rProgram 来利用多核并行性
Scaling up an rProgram
-
因为对齐特性确保了硬件效率，所以专注于通过构建正确的 rTile 形状来最大限度提高每个存储层的吞吐量
-
从初始 rTile 开始，逐步向最具 cost-efficient 的轴增大（即最大数据重用分数）--> 内存性能提高，直到达到计算极限或最大内存容量
-
自顶向下重复这个过程，直到构建出 rProgram
-
对于数据重用分数不变的算子（如 element-wise 操作），ROLLER 仅仅构建最顶层 rTile，然后直接从最底层存储层
-
算法：
Func ConstructProg(expr:TensorExpr, dev:Device): T = rTile(expr) Result = [] EnlargeTile(T, dev.MemLayer(0), rProg()) Func EnlargeTile(T:rTile, mem:MemLayer, P:rProg): if mem.IsLowestLayer(): # 最底层 Results.append(P) if (Results.Size() > TopK) Exit() # 找到 K 个 Return() for T_ in GetNextRTileShapes(T, mem): if Visited(T_): Return(); if MemFootprint(T_) > mem.Capacity(): # 数据占用超过存储容量 P.Add(mem, T) EnlargeTile(T, mem.Next(), P); else: if MemPerf(T_) > MaxComputePerf(T_.expr): # 内存加载数据吞吐量，超过峰值计算吞吐量 P.Add(mem, T_) EnlargeTile(T_, mem.Next(), P) EnlargeTile(T_, mem, P) Func GetNextRTileShapes(T:rTile, mem:MemLayer): alignedSizes = GetNextAlignedAxisSize(T, mem); SortedRTiles = OrderedMap(); for d in T.Dimensions: T_ = T.Replace(d, alignedSizes[d]); SortedRTiles.Insert({T_, DataReuseScore(T_)}) Return SortedRTiles
Scaling out an rProgram
- 同质性
- DNN 算子的计算模式
- 并行计算单元
- 单核 rProgram 复制到其他执行单元
- 将计算按底层 rTile 大小均匀划分
- 将这些 partition 均匀分配到所有执行单元上
- roller 更倾向于在同一执行单元上分配沿 reduce 轴拆分的 partition，已在更高的存储层中共享 reduce 结果
-
ROLLER 不会假设 rProgram 将独占所有的计算单元。系统可在 scale out 时显式控制 rProgram 的并行度
-
小算子与不规则张量形状
- 小算子
- scale-out 算法有利于有足够并行性的运算符，如 partition 数量明显大于执行单元数的运算符
- 对小算子来说，并行执行单元利用率低影响整体性能
- 若存在可并行的算子，可通过共同调度解决；否则，尝试沿着数据重用分数最小的轴缩减 rTile，以实现足够的并行性
- 不规则张量形状
- 大算子也可能因为拥有较小的不规则的张量形状，从而无法再满足对齐要求的前提下生成足够多的 rProgram
- 通过 axis fusion 来将其转换为规则形式
- 在所有张量中都存在且相邻（或都不存在）的维度
- 小算子
3.3 Efficient Evaluation of an rProgram
ROLLER 需要评估 rProgram 的性能，而非在真实硬件上评测端到端 rProgram。ROLLER 仅需评估相应 rTile 的性能，如 MemPerf 和 MaxComputePerf。
- HAL（hardware abstraction layer)
- HAL 将加速器建模为：带有分层存储层的多个并行执行单元
- 接口：Load，Compute，Store
- 执行单元抽象为 rTile Execution Unit（TEU）
- 通过 Compute 接口计算 data tiles
- 多个 TEU 组织为一个 group，共同进行 Load 和 Store 操作
- 存储层：
- 寄存器、共享内存、DRAM
- 特性：容量，传输长度，缓存行大小，内存 bank 数量
- Micro performance model：根据 HAL，轻松获得 rTile 的性能，从而获得 rProgram 的性能
- 对于给定的 rTile，如下数据可静态推断，用这些数据，计算数据重用分数，并检查 rTile 是否超过存储容量
MemFootprint：包括 padding 在内的内存占用MemTraffic：不同层间的存储流量
- 计算 rTile 的
MaxComputePerf- 通过一次性 profiling 测量峰值计算吞吐量：激进地扩大计算 tile 的大小来使得 TEU 饱和（该数据保留下，以便未来构造算法需要）
- 估计 rTile 的
MemPerf- 即，数据从低向高层存储层加载的性能
- 对于对齐的内存访问，加载规则数据块的延迟可简单地用总流量除以带宽来表示
- 对于所有 TEU 共享的存储层，平均划分内存带宽
- 对于小的访问大小，对每种设备类型用一次性的 profiling 来进行构造，并将结果进行保存。
- 对于给定的 rTile，如下数据可静态推断，用这些数据，计算数据重用分数，并检查 rTile 是否超过存储容量
4. Implementation
Roller on NVIDIA CUDA GPUs
- 内存架构
- global memory
- L2 cache
- L1 cache
- shared memory
- register
-
HAL
- L2：global memory + L2 cache
- L1：share memory
- L0：register
- （L1 cache 与 shared memory 共享空间，且用户程序无法控制，所以忽略）
-
内存带宽通过 micro-benchmark 测量
- 全局内存的传输长度：32 Bytes
- bank number：
- V100：32
- bank lenth：
- V100：4 Bytes
- K80：8 Bytes
- share memory 容量：48KB
- TEU：warp（32 线程）
- 在 HAL （即 SM）上 TEU Group 的大小是 warp 调度器的数量：4
-
SM 数量：
- V100：80
- K80：13
-
register：255 registers for V100
- nvcc 编译器会隐式地使用更多的寄存器（循环变量或其他目的）
- 设置为每线程 96 寄存器