发表时间: 2021-01 · arXiv:2101.08458
原文: https://arxiv.org/abs/2101.08458
文章作者/机构: Jian Weng, Animesh Jain, Jie Wang, Leyuan Wang, Yida Wang, Tony Nowatzki
所属机构: University of California, Los Angeles, USA; Amazon Web Services, USA
一句话结论 本文提出了一个名为UNIT的端到端编译流水线,通过为不同硬件平台的张量化指令建立统一的语义抽象,实现了深度学习模型张量操作的自动指令匹配与代码重写,在CPU和GPU上的推理性能均显著超越了官方供应商库。
要解决什么问题 深度学习中的混合精度计算(如使用低精度进行计算并用高精度进行累加)如果没有专门的硬件支持,数据类型转换的开销会完全抵消其性能优势。为此,主流硬件厂商推出了专门的张量化指令(如Intel VNNI、Nvidia Tensor Core和ARM DOT)。这些指令引入了一种名为“水平计算”的新模式,即先执行一系列元素级乘法,随后在同一个向量寄存器内部将结果直接累加成一个高精度元素。传统编译器的循环向量化器通常只针对最内层循环进行元素级的并行计算,遇到这种水平累加模式时,会错误地重排计算并生成冗余的归约代码,导致完全无法利用这些专用指令。现有的绕过方案要么依赖Intel oneDNN或Nvidia cuDNN等供应商库,这些库虽然高度优化但缺乏灵活性,难以支持新算子或算子融合;要么依赖开发者手动编写硬件内联函数,门槛极高且容易出错;而以往的专用编译器往往只针对单一平台的特定指令开发,当面对跨平台、多指令的场景时,开发成本难以承受。
怎么做的 解决这一卡点的核心思路是为所有张量化指令建立一个统一的语义抽象,从而将跨平台的分析和变换过程标准化。UNIT框架的运行机制由三个关键部件构成。首先是语义抽象部件,它使用张量领域特定语言将每条硬件指令视为一个小型的张量操作程序,统一描述其元素级运算和随后的水平归约行为,并将其解析为保留了表达式树和循环信息的张量操作数据结构。其次是适用性检测部件,它采用两步法自动判断指令能否用于目标操作。第一步检查计算同构性,通过递归遍历两棵表达式树来匹配数据类型和操作码,确保算术等价;第二步检查数组访问同构性,确保指令操作数张量中的每个元素只对应目标操作中的一个内存地址。其核心判定条件是:对于目标操作中的内存访问 $u$ 和指令中的内存访问 $v$ ,操作 $u$ 的索引表达式中被映射的循环变量集合 $S_0(u)$ 必须是指令 $v$ 索引表达式中循环变量集合 $S(v)$ 的子集,即满足 $$S_0(u) \subseteq S(v)$$ 如果满足该条件,说明目标操作加载的数据可以沿着缺失的循环变量进行广播以填满寄存器通道;如果不满足,则意味着一个寄存器通道会对应多个内存地址,映射不可行。最后是代码重写部件,分为三个阶段执行。第一阶段是循环重组,根据检测结果对目标循环进行切片,使其大小与指令维度一致,并重排到循环嵌套的最内层。第二阶段是指令替换,通过暴露待替换循环变量及其系数的统一编程接口,指导操作数进行向量化、广播或拼接,最终将最内层循环替换为目标指令。第三阶段是性能调优,针对未被张量化的外层循环,CPU端通过多线程实现粗粒度并行,并将少量数据并行循环重排展开到最内层归约循环下方,以避免写后读冒险;GPU端则采用外积风格的矩阵乘法累加来重用缓存子矩阵,并针对深度神经网络特有的形状引入维度融合以减少填充开销,以及对极深的通道维度进行归约拆分,在多线程并行计算后再通过共享内存同步完成最终归约。
效果如何 实验在三种硬件平台上展开:配备Intel VNNI指令的x86 CPU、配备Nvidia Tensor Core的V100 GPU,以及配备ARM DOT指令的ARM CPU。测试任务为深度学习模型的端到端推理,重点评估了最难优化的批处理大小为1的场景。对比基线包括代表极致优化的官方供应商库路线(Intel oneDNN、Nvidia cuDNN)、集成这些库的端到端框架路线(MXNet、TVM),以及代表手动优化路线的TVM手写内联函数方案。量化结果显示,UNIT在Intel CPU上比集成oneDNN的MXNet平均快1.3倍,比手动优化的TVM快1.18倍;在Nvidia GPU上比集成cuDNN的TVM平均快1.75倍,最高达2.2倍;在ARM CPU上比精心手动调优的TVM方案快1.13倍。此外,该方法展现了极强的扩展性,无需修改框架即可支持3D卷积等新算子,在Resnet18的3D变体上比oneDNN基线加速1.2倍。但作者也承认了该方法的局限性和失效场景:在CPU端,如果目标负载的输出形状无法被完美切片,处理循环余数时引入的条件分支会严重拖累性能;在GPU端,当遇到步进式数据访问时,数据局部性下降会导致性能不佳。此外,GPU上的展开度如果大于2会耗尽寄存器,且维度融合与归约拆分等专属优化本身也会引入数据重排和线程同步的开销,需要依赖调优器仔细权衡。
核心问题: 深度学习(DNN)依赖于密集的张量操作(如矩阵乘法和卷积),这些操作计算量巨大。混合精度计算(例如,使用低精度的fp16或int8进行计算,并用高精度的fp32或int32进行累加)是一种广泛采用的优化方法。然而,如果没有专门的硬件支持,数据类型转换的开销会抵消其性能优势。近年来,硬件供应商(如Intel、Nvidia、ARM)推出了专门用于混合精度张量操作的“张量化指令”(tensorized instructions),如Intel VNNI、Nvidia Tensor Core和ARM DOT。这些指令引入了一种新的计算模式——“水平计算”(horizontal computation),即在向量寄存器内部对多个低精度元素进行计算后累加成一个高精度元素。这种新模式给传统编译器带来了挑战,导致难以有效利用这些指令。
现有方案及其局限性:
研究目标与核心思想:
本文的目标是开发一种统一的编译方法,能够支持多种硬件平台上的不同张量化指令,以优化深度学习工作负载中的张量操作,并且该方法应具备良好的可扩展性。解决这一挑战的关键思想是为所有张量化指令建立一个统一的语义抽象(unified semantics abstraction),从而使分析和变换过程也能统一化。
本文贡献 (UNIT框架):
本文提出了一个名为UNIT的端到端编译流水线,以应对上述挑战。UNIT接收张量化指令的描述和深度学习模型作为输入,自动完成指令适用性检测、循环变换和代码重写,以高效利用这些专用指令。其主要贡献包括:
根据评估,UNIT能够支持包括Intel VNNI、Nvidia Tensor Core和ARM DOT在内的多种主流硬件指令。生成的端到端模型推理性能在x86 CPU上比Intel oneDNN快1.3倍,在Nvidia GPU上比Nvidia cuDNN快1.75倍,在ARM CPU上比精心手动调优的TVM方案快1.13倍。
深度学习计算成本高昂,而混合精度是降低计算和内存负担的有效技术。它通常使用低位宽数据类型(如fp16、int8)表示操作数,同时使用高位宽类型(如fp32、int32)进行累加以保证精度。然而,若没有硬件支持,数据类型转换的开销可能导致性能下降。实验表明,在没有Nvidia Tensor Core支持的情况下,使用fp16的cuDNN性能甚至低于fp32(如图1所示)。
为了解决此问题,主流硬件厂商推出了混合精度张量化指令,如Intel VNNI、ARM DOT和Nvidia Tensor Core,它们能够提供2到4倍的性能提升。这些指令引入了一种独特的“水平累加”模式:先执行一系列元素级乘法,然后将结果在同一个向量寄存器内进行累加。图2展示了Intel VNNI(4个int8元素的点积累加到int32)和Nvidia Tensor Core(16x16矩阵乘法)的语义。这种模式与传统的SIMD指令不同,后者通常只进行元素级的并行计算。
尽管张量化指令潜力巨大,但由于缺乏自动化的编译技术,其应用受到了限制。
编译器通常使用多级IR来在不同抽象层次上进行分析和变换。
tensor Op数据结构中。tensor Op被降级为张量IR,这是一种带有额外约束的命令式程序IR,例如所有循环都是从0开始、步长为1的规范形式。这些约束简化了分析和变换。UNIT在tensor Op级别进行分析,在张量IR级别进行变换。UNIT框架旨在自动地将混合精度张量操作映射到不同硬件平台的张量化指令上。其核心技术流程如图3所示,包括三个关键组件:语义抽象、适用性检测(Inspector)和代码重写(Rewriter)。
用Tensor DSL统一描述指令。为了统一编译流程并保持系统的可扩展性,UNIT使用通用的张量DSL来描述不同张量化指令的语义。每条指令被视为一个用DSL编写的小型张量操作程序。这种方法解决了指令语义的抽象问题。所有混合精度张量化指令都包含元素级运算和随后的水平归约。
a、b、c。计算行为由表达式d[i] = c[i] + sum(...)定义。循环i被标注为数据并行,而循环j是归约循环。+=),这是由于Tensor Core指令的数据类型不透明性所决定的。通过这种方式,指令的语义被解析为tensor Op数据结构,保留了表达式树、循环次数和数组缓冲区等信息,为后续的Inspector和Rewriter分析与变换提供了基础。
两步法确定指令适用性。Inspector组件采用两步法来判断一个张量化指令是否可以应用于某个张量操作。首先,通过检查表达式树的同构性来判断两者在算术上是否等价。然后,检查数据访问模式以确认操作数可以被正确准备,并指导Rewriter进行变换。
1. 计算同构性 (Compute Isomorphism):算法1描述了判断两个表达式树是否同构的过程。该算法递归地遍历两棵树,并匹配每对节点的 数据类型和操作码。图5(b).1展示了卷积操作和Intel VNNI指令(pbpdusd)的表达式树在拓扑结构、数据类型上完全相同,因此它们是算术同构的。此分析还会建立张量程序中的操作数到指令操作数(即抽象的寄存器)的映射关系,并确保一个寄存器不会对应多个数据源。
function INSPECT(a,b)
if a.type=b.type then
if isleaf(a)∧isleaf(b) then
if a is not bound then
bind[a]:=b
else if bind[a]!=b then
return False
end if
return True
else if isarith(a), isarith(b) then
cond:=a.opcode=b.opcode
cond:=cond∧Inspect(a.lhs, b.lhs)
cond:=cond∧Inspect(a.rhs, b.rhs)
return cond
end if
end if
return False
end function
算法 1:判断表达式树之间的同构性。a代表指令,b代表操作。
2. 数组访问同构性 (Array Access Isomorphism):在确定计算同构后,需要检查数据如何供给指令。此步骤的核心是确保指令操作数张量中的每个元素只对应于张量程序中的一个内存地址。UNIT通过枚举张量操作中哪些循环层级将被张量化(即映射到指令的循环),并检查这种映射是否可行。
u来自操作,v来自指令),操作u的索引表达式中被映射的循环变量集合$S_0(u)$必须是指令v索引表达式中循环变量集合$S(v)$的子集,即 $S_0(u) \subseteq S(v)$。三阶段代码变换。Rewriter组件分三个阶段进行代码变换:循环重组、张量化指令替换和性能调优。
c[x,y,k]被转换为16通道向量,a[x,y,rc]被相应地向量化和广播,而b[r,s,k,c]则被向量化、展开和拼接。threadIdx上并行化。所有分段完成后,通过线程同步在共享内存中进行最终的归约。UNIT通过扩展深度学习编译器Apache TVM【10】来实现,利用了其张量DSL、tensor Op、张量IR以及调优基础设施【11, Learning to optimize tensor programs. 2018. Advances in Neural Information Processing Systems】, 【23, Optimizing CNN model inference on CPUs. 2019. USENIX ATC 19】。
Inspector实现。通过分析TVM的ComputeOp数据结构实现,匹配表达式树并枚举循环变量的映射。它从张量的最内层维度向外层枚举,并贪婪地返回第一个符合条件的映射,以利用更好的数据局部性。
Rewriter实现。Rewriter按以下步骤执行:
tensorize pragma进行标注。tensorize pragma标注的循环体。调优策略实现:
NCHW[x]c数据布局和KCRS[y]k[x]c卷积核布局。batch size=1的情况。1. 端到端模型推理性能
2. 优化技术影响分析(消融实验)
实验选取了16个具有代表性的卷积层进行分析,其特征如表I所示。
表I:所选卷积层的特性。
3. 可扩展性
深度学习的发展促使硬件供应商增加了专门的张量化指令,这些指令引入了“水平归约”的计算模式,给通用编译带来了挑战,导致开发者不得不依赖手写核来获取高性能。本文介绍了UNIT,一个统一的编译流水线,它通过使用相同的IR来表示不同硬件平台的张量化指令,然后自动检测其适用性,变换循环嵌套,并用张量化指令重写循环体。UNIT实现了在Intel/ARM CPU和Nvidia GPU等多种硬件平台上的自动化张量化指令编译。实验表明,UNIT的性能显著优于供应商库(oneDNN加速1.3倍,cuDNN加速1.75倍)和手动调优方案(TVM ARM内联函数加速1.13倍)。
设置与复现指南。本附录描述了如何设置UNIT编译环境并运行第VI节中讨论的实验。内容包括:通过Docker设置实验环境;运行端到端模型推理实验(图8、9、12);运行调优策略效果的实验(图10、11);以及运行3D卷积实验(图13)。
实验环境与依赖。实验在Amazon EC2实例上进行(c5.12xlarge用于Intel VNNI,p3.2xlarge用于Nvidia TensorCore,m6g.8xlarge用于ARM VDOT)。需要约32GB磁盘空间。提供了Dockerfile来搭建环境,并有脚本自动运行实验和绘图。
软硬件依赖。硬件依赖于特定CPU/GPU架构(如Intel Cascade Lake、Nvidia Volta、ARMv8.2-dotprod)。软件依赖项通过Docker自动安装。
安装与实验流程。提供了详细的步骤,指导用户解压文件,在不同平台(GPU、CPU、ARM)上构建Docker镜像,进入容器,并执行脚本来复现论文中的实验结果。脚本会自动运行实验并生成PDF格式的图表。
评估与预期结果。指出最终生成的PDF文件应与论文中的图表进行对比。同时提到,由于在Docker虚拟机中运行,性能可能存在波动。ARM的结果是基于旧版TVM生成的,新版性能有提升。