UNIT: Unifying Tensorized Instruction Compilation

文章作者/机构: Jian Weng, Animesh Jain, Jie Wang, Leyuan Wang, Yida Wang, Tony Nowatzki
所属机构: University of California, Los Angeles, USA; Amazon Web Services, USA

A1 主要贡献

核心问题: 深度学习（DNN）依赖于密集的张量操作（如矩阵乘法和卷积），这些操作计算量巨大。混合精度计算（例如，使用低精度的fp16或int8进行计算，并用高精度的fp32或int32进行累加）是一种广泛采用的优化方法。然而，如果没有专门的硬件支持，数据类型转换的开销会抵消其性能优势。近年来，硬件供应商（如Intel、Nvidia、ARM）推出了专门用于混合精度张量操作的“张量化指令”（tensorized instructions），如Intel VNNI、Nvidia Tensor Core和ARM DOT。这些指令引入了一种新的计算模式——“水平计算”（horizontal computation），即在向量寄存器内部对多个低精度元素进行计算后累加成一个高精度元素。这种新模式给传统编译器带来了挑战，导致难以有效利用这些指令。

现有方案及其局限性:
1. 供应商库 (Vendor Libraries): 如Intel oneDNN和Nvidia cuDNN，为预定义的计算核提供了高度优化的性能。但缺点是缺乏灵活性，难以支持新的工作负载或进行更深层次的优化（如算子融合）。
2. 手动编写硬件内联函数 (Intrinsics): 这种方式对普通开发者门槛极高，容易出错且效率低下。
3. 专用编译器: 以往的一些工作为每种特定的张量化指令开发一个编译器，例如PolyDL【36, PolyDL: Polyhedral optimizations for creation of high performance DL primitives. 2020. arXiv】和Bhaskaracharya等人【35, Automatic kernel generation for Volta tensor cores. 2020. arXiv】的工作。当需要支持跨平台的多种指令时，这种方法需要投入过多的开发精力。

研究目标与核心思想:
本文的目标是开发一种统一的编译方法，能够支持多种硬件平台上的不同张量化指令，以优化深度学习工作负载中的张量操作，并且该方法应具备良好的可扩展性。解决这一挑战的关键思想是为所有张量化指令建立一个统一的语义抽象（unified semantics abstraction），从而使分析和变换过程也能统一化。

本文贡献 (UNIT框架):
本文提出了一个名为UNIT的端到端编译流水线，以应对上述挑战。UNIT接收张量化指令的描述和深度学习模型作为输入，自动完成指令适用性检测、循环变换和代码重写，以高效利用这些专用指令。其主要贡献包括：
* 统一的张量化指令抽象: 提出了一种使用张量领域特定语言（Tensor DSL）来统一描述不同硬件平台张量化指令语义的方法。
* 指令适用性检测算法: 设计了一种算法，能够自动检测给定的张量化指令是否以及如何应用于目标张量操作。
* 代码重写与调优机制: 开发了一套代码重写和自动调优机制，能够寻找最优的循环变换方式，将张量化指令嵌入到张量操作中以实现高性能。

根据评估，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倍。

A3 背景知识

混合精度张量化指令

深度学习计算成本高昂，而混合精度是降低计算和内存负担的有效技术。它通常使用低位宽数据类型（如fp16、int8）表示操作数，同时使用高位宽类型（如fp32、int32）进行累加以保证精度。然而，若没有硬件支持，数据类型转换的开销可能导致性能下降。实验表明，在没有Nvidia Tensor Core支持的情况下，使用fp16的cuDNN性能甚至低于fp32（如图1所示）。

图1：在Nvidia V100-SXM2上，无混合精度指令支持时fp32与fp16的性能对比。

为了解决此问题，主流硬件厂商推出了混合精度张量化指令，如Intel VNNI、ARM DOT和Nvidia Tensor Core，它们能够提供2到4倍的性能提升。这些指令引入了一种独特的“水平累加”模式：先执行一系列元素级乘法，然后将结果在同一个向量寄存器内进行累加。图2展示了Intel VNNI（4个int8元素的点积累加到int32）和Nvidia Tensor Core（16x16矩阵乘法）的语义。这种模式与传统的SIMD指令不同，后者通常只进行元素级的并行计算。

图2：Intel VNNI和Nvidia Tensor Core的语义。旁边的文本是对应的LLVM内联函数名称。

现有解决方案的局限性

尽管张量化指令潜力巨大，但由于缺乏自动化的编译技术，其应用受到了限制。
* 供应商库（如Intel oneDNN, Nvidia cuDNN）虽然性能高，但不够灵活，无法应对预定义操作之外的新需求。
* 传统循环向量化器（如GCC, LLVM中的）主要针对最内层循环的并行化，难以处理“水平累加”这种复杂的计算模式。它们通常会重排计算并生成收尾（epilogue）归约代码，从而无法利用张量化指令。
* 近期的编译工作，如PolyDL【36】和Bhaskaracharya等人【35】的工作，分别针对Intel VNNI和Nvidia Tensor Core，但它们局限于特定平台和指令，缺乏通用性。

多级中间表示 (IR)

编译器通常使用多级IR来在不同抽象层次上进行分析和变换。
1. 图级别IR (Graph-Level IR): 深度学习编译器如TVM【10, TVM: An automated end-to-end optimizing compiler for deep learning. 2018. 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)】、Glow【34, Glow: Graph lowering compiler techniques for neural networks. 2018. CoRR】和XLA【43, Xla - tensorflow, compiled. 2017】使用图级别IR将模型表示为算子的有向无环图（DAG）。这有助于进行张量间优化，如算子融合和数据布局选择。UNIT依赖于此级别的张量填充（padding）操作。
2. 张量领域特定语言 (Tensor DSL): Halide【31, Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines. 2013. PLDI '13】、TVM【10】等DSL允许开发者高效地描述张量计算，并提供了接口来对循环进行拆分、重排等变换以进行性能调优。所有信息被存储在tensor Op数据结构中。
3. 张量IR (Tensor IR): tensor Op被降级为张量IR，这是一种带有额外约束的命令式程序IR，例如所有循环都是从0开始、步长为1的规范形式。这些约束简化了分析和变换。UNIT在tensor Op级别进行分析，在张量IR级别进行变换。
4. 低级IR (Low-Level IR): 张量IR最终被降级为通用的低级IR（如LLVM IR），为生成汇编代码做准备。

A2 方法细节

UNIT框架旨在自动地将混合精度张量操作映射到不同硬件平台的张量化指令上。其核心技术流程如图3所示，包括三个关键组件：语义抽象、适用性检测（Inspector）和代码重写（Rewriter）。

A. 语义抽象 - 张量DSL

用Tensor DSL统一描述指令。为了统一编译流程并保持系统的可扩展性，UNIT使用通用的张量DSL来描述不同张量化指令的语义。每条指令被视为一个用DSL编写的小型张量操作程序。这种方法解决了指令语义的抽象问题。所有混合精度张量化指令都包含元素级运算和随后的水平归约。
* Intel VNNI指令的DSL描述：如图4(a)所示，VNNI的三个源操作数（两个int8/uint8，一个int32）被定义为张量a、b、c。计算行为由表达式d[i] = c[i] + sum(...)定义。循环i被标注为数据并行，而循环j是归约循环。
* ARM DOT指令的DSL描述：如图4(b)所示，其描述与VNNI类似，但通道数和数据类型不同。
* Nvidia Tensor Core指令的DSL描述：如图4(c)所示，它执行一个16x16的方阵乘法。与前两者不同的是，其累加器寄存器必须与加法寄存器相同（+=），这是由于Tensor Core指令的数据类型不透明性所决定的。

通过这种方式，指令的语义被解析为tensor Op数据结构，保留了表达式树、循环次数和数组缓冲区等信息，为后续的Inspector和Rewriter分析与变换提供了基础。

B. 适用性检测 - Inspector

两步法确定指令适用性。Inspector组件采用两步法来判断一个张量化指令是否可以应用于某个张量操作。首先，通过检查表达式树的同构性来判断两者在算术上是否等价。然后，检查数据访问模式以确认操作数可以被正确准备，并指导Rewriter进行变换。
* (a). 算法描述

// 张量DS中的卷积 a,b = tensor((H,W,C), u8),tensor((R,S,K,C),i8) k,rc = loop_axis(0,K), reduce_axis(0,C) x,y = loop_axis(0,H-R+1), loop_axis(0,W-S+1) r,s = reduce_axis(0,R), reduce_axis(0,S) c[x,y,k]+= i32(a[x+r,y+s,rc])*i32(b[r,s,k,rc]) （b）.1 算术同构

**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代表操作。

（b）.2 数据访问同构

2. 数组访问同构性 (Array Access Isomorphism)：在确定计算同构后，需要检查数据如何供给指令。此步骤的核心是确保指令操作数张量中的每个元素只对应于张量程序中的一个内存地址。UNIT通过枚举张量操作中哪些循环层级将被张量化（即映射到指令的循环），并检查这种映射是否可行。
* 可行性检查：一个映射被认为是可行的，需要满足以下条件：对于每一对对应的内存操作（u来自操作，v来自指令），操作u的索引表达式中被映射的循环变量集合$S_0(u)$必须是指令v索引表达式中循环变量集合$S(v)$的子集，即 $S_0(u) \subseteq S(v)$。
* 条件解读：如果满足该子集关系，意味着张量操作加载的数据需要沿着$S(v)$中存在但$S_0(u)$中不存在的循环变量进行广播，以填满所有寄存器通道。如果不满足，则意味着一个寄存器通道将对应多个内存地址，这在代码生成中是不现实的。
* 如图5(b).2所示，若存在多种可行的映射，它们将构成代码调优空间的一部分。确定的映射将指导后续的循环变换和代码生成。

C. 代码变换 - Rewriter

三阶段代码变换。Rewriter组件分三个阶段进行代码变换：循环重组、张量化指令替换和性能调优。
1. 循环重组 (Loop Reorganization)：根据Inspector选择的待张量化循环层级，Rewriter首先对这些循环进行切片（tiling），切片的大小与指令的循环维度相对应。然后，将这些切片后的循环重排到循环嵌套的最内层，使得最内层循环的语义与指令的语义完全一致。如图5(c)所示，这一过程通过张量DSL提供的接口可以轻松实现。
2. 张量化指令替换 (Tensorized Instruction Replacement)：在识别出待替换的代码区域后，代码生成器需要为指令准备好每个操作数。由于不同平台的执行模型和汇编格式各异，完全自动化操作数准备是困难的。因此，UNIT定义了一个统一的编程接口，允许编译器开发者手动指定操作数生成规则。该接口暴露了待替换的循环变量及其在索引表达式中的系数。例如，在图5(c)中，通过分析$k_i$和$c_i$的步长和循环次数，c[x,y,k]被转换为16通道向量，a[x,y,rc]被相应地向量化和广播，而b[r,s,k,c]则被向量化、展开和拼接。
3. 调优器 (Tuner)：未参与指令重写的其他循环层级可以被重组以调优性能。UNIT为CPU和GPU设计了不同的调优策略，通用理念是同时利用细粒度和粗粒度的并行性。
* CPU调优：
* 粗粒度并行：通过将数据并行循环分配到多个线程来实现。
* 指令级并行：为了避免归约循环中由循环携带依赖（loop-carried dependence）引起的RAW（Read-After-Write）冒险，将少量数据并行循环重排并展开到最内层归约循环的下方。
* 调优空间：展开的程度和并行化的线程数是两个调优维度。UNIT通过枚举和性能剖析来寻找最佳组合，以平衡指令级并行、I-cache命中率、核心利用率和上下文切换开销。
* GPU调优：
* 粗粒度并行：通过将数据并行循环分配到流式多处理器（SM）上实现。
* 数据重用：在GPU上，数据重用由软件显式管理。如图6所示，采用外积（outer-product）风格的矩阵乘法累加，以重用缓存的子矩阵。
* 图6：累加一个p×p的“方形窗口”可以避免循环携带的数据依赖，并重用缓存的子矩阵。
* DNN核专属优化：
* 维度融合 (Dimension fusion)：对于宽高较小的层，将这两个维度融合成一个，以节省冗余的填充开销。
* 归约拆分 (Split reduction)：对于通道数很深（即归约循环次数大）的层，将归约循环拆分，并将每个分段在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按以下步骤执行：
1. 根据Inspector的分析结果，通过切片和重排来重组待张量化的循环，并用tensorize pragma进行标注。
2. 根据第III-C节讨论的策略，重组外部循环以进行性能调优。
3. 将操作后的循环嵌套降级到张量IR，并用目标指令替换被tensorize pragma标注的循环体。

调优策略实现：
* CPU调优：如图7所示，代码通过定义两个“断点”（breaking points）来划分循环。第一个断点前的循环被融合并并行化；两个断点之间的循环串行执行；第二个断点后的循环被重排到最内层并展开。
* 图7：CPU调优的代码示意图。
* GPU调优：在粗粒度和细粒度并行之间存在权衡。实验发现，展开度（图6中的p）大于2可能会耗尽寄存器，因此使用p=2。对于DNN的专属优化（维度融合、归约拆分），由于它们可能引入自身的开销（数据重排、线程同步），调优器会枚举相关参数（如归约并行度、是否融合维度）并进行性能剖析，以确定最佳的变换组合。

A4 实验环境与结果

实验环境

硬件平台:
- Intel x86 CPU: Amazon EC2 C5.12xlarge实例，配备24核Intel Xeon Platinum 8275CL CPU (Cascade Lake)。
- ARM CPU: Amazon EC2 M6g.8xlarge实例，配备32核AWS Graviton2 CPU (ARM Cortex-A72)。
- Nvidia GPU: Amazon EC2 P3.2xlarge实例，配备Nvidia Tesla V100 SXM2 GPU。
软件配置:
- 所有程序通过Apache TVM实现，并使用LLVM-10作为后端生成代码。GPU部分使用CUDA-10.0。
- 基线:
  - 供应商库: Intel oneDNN v1.6.1 和 Nvidia cuDNN 7.6.5。
  - 端到端基线: 集成oneDNN的MXNet (CPU) 和集成cuDNN的TVM (GPU)。
  - 手动优化基线: TVM中为Intel VNNI和ARM DOT指令手动编写内联函数的现有解决方案。
数据集与模型:
- DNN模型: 所有模型均来自MXNet Model Zoo，并转换为TVM的图IR (Relay)。
- 数据布局: 采用NCHW[x]c数据布局和KCRS[y]k[x]c卷积核布局。
- 评测场景: 重点评估推理场景中关键但难以优化的batch size=1的情况。

实验结果

1. 端到端模型推理性能
* Intel x86 (VNNI): 如图8所示，与集成Intel oneDNN的MXNet相比，UNIT实现了1.3倍的平均加速；与手动优化的TVM相比，实现了1.18倍的加速。尽管oneDNN在某些模型（如resnet50）上经过了深度优化，UNIT在整体上表现更优。
图8：由Intel VNNI加速的量化网络推理（bs=1）。
* Nvidia GPU (Tensor Core): 如图9所示，与集成cuDNN的TVM相比，UNIT始终表现更佳，平均加速比为1.75倍，最高可达2.2倍。
图9：由Tensor Core加速的混合精度网络推理（bs=1）。

2. 优化技术影响分析（消融实验）
实验选取了16个具有代表性的卷积层进行分析，其特征如表I所示。
* Intel x86: 如图10所示，性能提升主要来自并行化（Parallel）和展开（Unroll）的结合。自动调优（Tune）带来的额外增益较小，因为超过95%的核在前8个调优配置中就达到了最优性能。在#1和#4两个负载上表现不佳，因为其输出形状无法被完美切片，导致循环余数处理引入的if分支影响了性能。
图10：代码空间探索的性能影响。
* Nvidia GPU: 如图11所示，归约拆分（+SplitK）带来了最大的性能提升，因为它显著增加了并行度，使Tensor Core保持繁忙。维度融合（+FuseDim）对某些形状的核也有帮助。与CPU类似，调优的额外增益不大。在#1和#15两个负载上表现不佳，因为步进式（strided）数据访问降低了数据局部性。尽管存在这些不利情况，但UNIT的通用优化方法使其在整体上优于供应商库。
图11：代码空间探索的性能影响。

表I：所选卷积层的特性。

3. 可扩展性
* 对新硬件平台的支持: 将UNIT应用于支持ARM DOT指令的ARM CPU。如图12所示，UNIT的性能优于基线TVM-NEON和手动优化的TVM-Manual，证明了其只需在DSL中描述新指令语义即可轻松扩展到新平台的能力。
图12：ARM在模型推理上的性能。
* 对新张量操作的支持: 将UNIT应用于3D卷积，映射到Intel VNNI指令。如图13所示，在Resnet18的3D变体上，UNIT相比oneDNN基线取得了平均1.2倍的加速，表明其无需修改框架即可扩展到新的张量操作。
图13：res18-3d上各层的性能。

A5 结论

深度学习的发展促使硬件供应商增加了专门的张量化指令，这些指令引入了“水平归约”的计算模式，给通用编译带来了挑战，导致开发者不得不依赖手写核来获取高性能。本文介绍了UNIT，一个统一的编译流水线，它通过使用相同的IR来表示不同硬件平台的张量化指令，然后自动检测其适用性，变换循环嵌套，并用张量化指令重写循环体。UNIT实现了在Intel/ARM CPU和Nvidia GPU等多种硬件平台上的自动化张量化指令编译。实验表明，UNIT的性能显著优于供应商库（oneDNN加速1.3倍，cuDNN加速1.75倍）和手动调优方案（TVM ARM内联函数加速1.13倍）。

A6 附录

设置与复现指南。本附录描述了如何设置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生成的，新版性能有提升。