一种与几何无关的矩阵乘法指令集架构扩展

《Future Generation Computer Systems》：A geometry-agnostic matrix-multiplication instruction Set Architecture extension

　　

摘要
本文提出了矩阵瓦片扩展（MTE），一种新颖的矩阵指令集架构（ISA），旨在通过几何无关的矩阵乘法指令来补充新兴的向量长度无关架构。受现代矩阵扩展启发，MTE 的独特之处在于将应用程序矩阵形状与底层矩阵乘累加加速器设计完全解耦，同时与基础向量 ISA 无缝集成。MTE 通过以下方式实现了极低的实现开销：(i) 复用向量寄存器堆进行矩阵存储，(ii) 仅引入少量额外指令，(iii) 在单个 64 位控制和状态寄存器中维持其执行状态。我们将 MTE 描述为一种 RISC-V 自定义扩展，并在机器学习卷积和矩阵乘法工作负载上评估了我们的提案。为了评估现有矩阵 ISA 的局限性以及我们方法的益处，我们将对最先进矩阵 ISA 瓶颈的经验性特征分析与 MTE 的轨迹驱动模拟方法相结合。我们的结果表明，MTE 相比当代矩阵 ISA 实现了高达 1.35 倍的加速，这主要是通过利用复用向量寄存器堆所启用的更大架构状态，以及消除固定加速器几何形状所施加的约束来实现的。

引言
通用矩阵乘法（GEMMs）在高性能计算和深度学习工作负载中无处不在 [1]。为满足现代应用日益增长的算术吞吐量需求，主要硬件供应商现在将矩阵乘法加速器集成到通用 CPU 架构中 [2]，[3]，[4]，[5]。矩阵 ISA 提供了显著的好处，但需要仔细的架构支持来管理矩阵操作数并提供简化的软件/硬件接口。常见方法包括向量寄存器分组 [6]、专用瓦片寄存器 [2]，或能够存储 $n^2$ 个元素的扩展寄存器，其中 $n$ 是向量/SIMD 宽度 [7]。虽然有效，但此类技术在基于短向量指令集架构（ISA）[3]，[8]，[9] 及其原始设计之外的灵活性有限，且仅适用于近似方形的工作负载。

本文提出了矩阵瓦片扩展（MTE），一种新颖的矩阵 ISA 扩展，适用于向量长度无关架构，通过几何无关的软件编程模型，为短向量和长向量处理器均提供加速器设计灵活性。MTE 可以 (i) 在整个 GEMM 循环嵌套 (M, N 和 K) 中向量化矩阵乘法；(ii) 在向量和矩阵计算期间充分利用向量架构状态容量；以及 (iii) 将软件几何形状与底层微架构解耦。通过利用这三个特性，MTE 克服了现有向量 [8]，[10]，[11] 和矩阵 [2]，[4] ISA 在处理来自相关人工智能（AI）模型的异构 GEMM 工作负载时的性能瓶颈。具体来说，我们在本文中做出了以下贡献：
• 我们确定了与现代矩阵 ISA 相关的关键性能瓶颈，涉及 (i) 矩阵存储空间利用率低，(ii) 寄存器内矩阵布局，以及 (iii) 暴露给软件的寄存器数量受限。
• 我们提出了 MTE，这是首个适用于短向量和长向量的几何无关 ISA，实现了软件在不同实现之间的可移植性。我们描述了两种 MTE 实现，一种是基于长向量处理器的精简扩展，另一种利用了脉动阵列。
• 我们评估了 MTE，考虑了属于计算机视觉深度神经网络 [12]，[13]，[14]，[15]，[16] 的 75 个卷积工作负载和用于自然语言处理的变压器模型 [17]，[18] 的 18 个矩阵乘法工作负载。

我们的评估方法包括两种基于具有 8192 位和 16,384 位向量寄存器的向量 ISA 的方法，以及两个最近的矩阵 ISA 提案 [2]，[4]。我们的评估表明，当在高性能架构上实现时，MTE 比之前的方法提供了更好的性能。特别是，由于软件利用了更多架构可见寄存器，MTE 相比最先进的方法获得了 1.35 倍的几何平均加速比。

章节片段
背景与动机
本节描述了 GEMM 操作的标准软件接口（第 2.1 节），介绍了现代向量架构（第 2.2 节）、最先进的 CPU 矩阵 ISA（第 2.3 节）及其不足之处（第 2.4 节）。

矩阵瓦片扩展
MTE 利用外积执行模型来最大化矩阵乘法期间的操作数重用，如图 3 所示。它包括将特定形状 (m, n) 的输出 C 矩阵瓦片存储到架构状态中，并根据方程 $C_{m,n}=\sum_{k=0}^{K-1} A_{m,k} B_{k,n}$ 并行更新每个元素。对于沿归约维度 K 的每一步，矩阵 A 的一个切片和矩阵 B 的相应切片被加载到架构状态中。然后硬件发出一组 FMAs 来……

MTE 指令的架构支持
本节描述了 MTE 所需的架构支持。第 4.1 节描述了 MTE MMA 指令的两种可能的微架构实现。第 4.2 节描述了 MTE 内存指令所需的支持，而第 4.3 节讨论了其他方面。

系统架构
表 4、5 和 6 展示了我们在实验活动中考虑的系统、向量处理单元和脉动阵列架构。系统和脉动阵列参数与 Intel Xeon Platinum 8480+ 处理器 [34] 核心相匹配，这是实现 AMX 扩展的最新 x86 服务器级处理器。向量处理单元 (VPU) 参数类似于单个 NEC SX-Aurora 20B 向量加速器核心 [44]，增加了一个向量单元，总计四个。每个单元具有 2048 位……

评估
我们通过运行提议的工作负载（第 5.3 节变压器工作负载，5.4 节卷积工作负载）来评估第 5.5 节的技术，并讨论其性能（第 6.1 节）、能耗（第 6.2 节）和已退休指令数（第 6.3 节）。

每种策略都实现为针对不同张量内存布局的 oneDNN 操作符。在矩阵乘法期间，我们的基本函数使用行优先张量布局。在卷积上，我们的基本函数使用分块内存布局 [21]，[29]，……

相关工作
本节详细介绍了相关现代矩阵 ISA 的主要方面，并在表 10 中概述了它们的主要差异。

OpenPower ISA v3.1 (OpenPower) [6] 引入了混合精度、单精度、双精度和整数矩阵的外积指令，使用 512 位累加器，并在累加器与 128 位向量寄存器之间进行寄存器到寄存器的传输。新外积指令从向量寄存器获取 A 和 B 操作数，从更大的累加器获取 C 操作数。

结论
本文展示了现有向量和矩阵 ISA 在处理基于 GEMM 的工作负载时的局限性。本文将最先进方法在处理高瘦或小矩阵时的缺陷归因于向量寄存器堆的利用率不足，并提出了一种精简的矩阵瓦片扩展 MTE 来解决这一问题。MTE 将向量寄存器重新解释为由新颖的几何无关矩阵 ISA 操作的矩阵瓦片操作数。MTE 解耦了矩阵……

CRediT 作者贡献声明
Alexandre de Limas Santana: 写作 – 审阅与编辑，写作 – 初稿，可视化，验证，软件，方法论，调查，概念化。Adrià Armejach: 写作 – 审阅与编辑，写作 – 初稿，验证，监督，方法论。Francesc Martinez: 写作 – 初稿，可视化，方法论。Erich Focht: 资源。Marc Casas: 写作 – 审阅与编辑，写作 – 初稿，监督，方法论。

利益冲突声明
作者声明以下财务利益/个人关系，这可能被视为潜在的利益冲突：Marc Casas 报告国际商业机器公司（IBM）提供了财务支持。如果有其他作者，他们声明没有已知的竞争性财务利益或个人关系可能会影响本文报告的工作。

致谢
本工作获得了 EuroHPC-JU 的资助，拨款编号为 101034126，并得到 Horizon2020 项目的支持。本工作是西班牙科学、创新与大学部 MCIU/AEI/10.13039/501100011033 资助项目 PID2023-146511NB-I00 的一部分，并得到欧盟 ERDF 的支持。本工作已获得“计算的未来，巴塞罗那超级计算中心与 IBM 倡议”（2023 年）的资助。

Alexandre de Limas Santana | Adrià Armejach | Francesc Martinez | Erich Focht | Marc Casas
巴塞罗那超级计算中心，Pla?a Eusebi Güell, 1-3, Barcelona, 08034, Catalunya, Spain