对稀疏矩阵运算的两种优化方式

背景卷积神经网络CNN广泛应用于移动端视觉任务GEMM 是其推理的性能瓶颈脉动阵列SA通过局部寄存器通信高效加速 GEMM被广泛应用于 TPU 等商用产品但传统架构仍有优化空间。面临问题CNN 模型中普遍存在稀疏性但利用稀疏性面临四大挑战随机稀疏难以利用、结构化稀疏影响精度、索引存储开销大、SA 寄存器比率失衡。本文提出的 STA 架构一次性解决了这些问题。问题描述随机稀疏难以利用零元素位置无规律需额外存储索引负载不均衡导致硬件利用率低下结构化稀疏影响精度将零值按固定块组织虽便于硬件实现但块尺寸较大时会显著损害 CNN 精度索引存储开销大传统方法需为每个非零元素存储索引在 INT8 权值场景下开销尤为严重SA 寄存器比率失衡每个 MAC 单元需要两个 INT8 寄存器 一个 INT32 累加器比例存在优化空间1. 引言本文两大贡献将标量 PE 泛化为 Tensor-PE 衍生出 STA 架构引入 DBB 稀疏格式实现 STA-DBB在面积效率和功耗上均大幅领先 SA 基线同时完全向后兼容密集模型。卷积层通常通过将二维卷积降维为 GEMM 核来执行。SA 由一列 MAC 处理单元PE构成各单元之间仅通过局部的寄存器到寄存器通信传递操作数和结果使其非常高效且易于扩展不产生时序问题。这些优势使其被部署在 Google TPU 等商用产品中。本文针对移动端 CNN 推理的窄 INT8 操作数描述了两项对传统 SA 架构的重大改进第一STA将传统的标量 PE 泛化为Tensor-PE衍生出Systolic Tensor ArraySTA微架构家族。STA 通过增加 PE 内操作数复用和 datapath 效率在同吞吐量下与时钟门控 SA 相比电路面积减少最高2.08 倍功耗降低1.36 倍。第二STA-DBB扩展 STA 以支持Density-Bound BlockDBB新型块稀疏数据格式。该变体在处理专门训练的 DBB 稀疏模型时相较于 SA 基线在同吞吐量下面积效率提升3.14 倍功耗效率提升1.97 倍同时完全向后兼容密集模型。2. 背景与相关工作2.1 面向 CNN 推理的脉动阵列INT8 微架构的 datapath 成本相对较低数据移动成本占比更高本文以 TPU-like SA 为基线通过时钟门控和输出驻留数据流进行优化。本文以 TPU-like SA 为起点进行了修改1对零操作数进行时钟门控2采用输出驻留数据流——让较大的 INT32 累加器保持不动而移动较小的 INT8 操作数。2.2 CNN 的稀疏矩阵DBB 是随机稀疏与结构化块稀疏之间的完美折中对每个块的非零元素数量设定上限在相同 NNZ 下获得更高精度同时硬件计算量已知保证高利用率。CNN 层在权值数据和激活数据中均存在稀疏性。随机稀疏图 1a零的位置没有约束很难利用结构化块稀疏图 1b硬件友好但块尺寸较大时会损害 CNN 精度本文提出DBB图 1c对每个块的非零NNZ元素数量设定上限。例如由 8 个 8×1 块组成每个块最多 3 个非零值NNZ≤3。这在相同 NNZ 下能获得更高 CNN 精度同时硬件计算量已知保证高利用率。DBB 采用位掩码压缩可实现权重内存占用减少 37.5%。Fig 1. 三种稀疏矩阵格式对比(a) 随机稀疏(b) 块稀疏4×2 块© DBB8×1 块。填充元素表示非零值。2.3 稀疏矩阵乘法加速器索引方法存储开销巨大本文在 SA 架构基础上另辟蹊径采用 Tensor-PE 新型结构化稀疏方案而非随机稀疏。利用零操作数的方法包括时钟门控和索引存储例如 EIE、SCNN。然而索引方法开销巨大。本文在 SA 架构基础上利用Tensor-PE 新型结构化稀疏方案而非随机稀疏开创了一条新路径。3. 密集 GEMM 加速器3.1 传统脉动阵列SASA 通过局部寄存器通信实现高效 GEMM但存在寄存器/MAC 比率失衡的优化空间基线 SA 采用 INT8 操作数寄存器和 INT32 累加器。Fig 2. SA 微architectures将每个标量 PE 泛化为执行矩阵乘法的 tensor PE。NotationA×B×C_M×N 表示 M×N 个 A×B×C tensor PE 的阵列。图 2a 展示了经典的 SA被广泛应用于产品中。目标应用为移动视觉场景采用 INT8 操作数寄存器REG和 INT32 累加器寄存器ACC。传统 SA 仍有很大改进空间——放松操作数寄存器REG与 MAC 的比率。3.2 脉动张量阵列STASTA 将多个标量 PE 融合为一个 Tensor-PE每个 Tensor-PE 内执行向量点积经典 SA 是 STA 的特例ABC1通过减少操作数缓冲区和累加器缓冲区实现显著效率提升。Fig 3. Example 2×2×2_2×2 STA 数据流。tensor PE 执行 4×4 × 4×4 矩阵乘法。图 2b 展示了传统 SA图 2a 展示了Systolic Tensor ArraySTA——通过将一组标量 PE 融合为单个tensor PE来实现泛化。每个 STA 架构包含 M×N 个 tensor PE每个 tensor PE 由 A×C 个 MAC 子阵列构成每个 MAC 对 B 对操作数执行点积运算表示为A×B×CM×N A \times B \times C_{M \times N}A×B×CM×N​图 2c 展示了一个 2×2 的 tensor PE 阵列每个 tensor PE 有 2×2 的 datapath包含 4 个操作数对点积单元DP4。经典 SA 是 STA 的特例ABC1即 1×1×1_M×N。通过将每个 MAC 的操作数缓冲区减少2 倍累加器缓冲区减少4 倍实现了显著的效率提升。4. CNN 稀疏 GEMM 加速器4.1 Density-Bound BlockDBBDBB 对每个块的非零元素数量设定上限是随机稀疏与块稀疏之间的完美折中精度损失仅 0.1%~1.1%内存减少 37.5%硬件高利用率。【请在此处插入 Fig 1 图片liu2020_images/fig1_page2.jpeg】Fig 1©. DBB8×1 块稀疏矩阵格式对每个块的 NNZ 设定上限。DBB 矩阵格式的核心理念很简单对每个块的非零NNZ元素数量设定上限。这与传统的块稀疏格式不同——在传统格式中每个块要么完全不受约束要么全为零。DBB 是随机稀疏与块稀疏之间的完美折中在相同 NNZ 下能获得更高的 CNN 精度。DBB 采用简单的位掩码压缩来编码每个 8 元素块每个块 1 字节开销外加 4 字节非零数据可实现权重内存占用减少 37.5%。4.2 面向 DBB 的 STASTA-DBBSTA-DBB 利用 DBB 上限削减物理 MAC 单元节省 50%用 MUX 替代乘法器成本远低于 8 位乘法器保持一半吞吐量密集 GEMM 向后兼容。将 STA 架构扩展以支持 DBB 稀疏权值矩阵。当 DBB 块大小为 8、NNZ≤4 时每个 8 输入点积单元DP8只需 4 个 MAC 单元而非 8 个实现物理 MAC 减少 50%。该方案在 MAC 前加一个多路复用器MUX根据非零权值的索引选择对应的激活值。在每个 SDP4 中用两个 8 位 4:1 MUX 换取两个 8 位乘法器——MUX 的成本远低于 8 位乘法器。该阵列在每个时钟周期执行16 个有效 MAC但仅使用8 个物理乘法器。重要的是该架构仍以一半吞吐量支持传统密集 GEMM完全向后兼容。5. 结论本文通过将标量 PE 泛化为 Tensor-PESTA增加 PE 内操作数复用显著减少电路面积和功耗进一步扩展的 STA-DBB 在处理 DBB 稀疏模型时效率和兼容性兼得完全向后兼容传统密集模型。文章信息标题Systolic Tensor Array: An Efficient Structured-Sparse GEMM Accelerator for Mobile CNN Inference作者Zhi-Gang Liu, Paul N. Whatmough, Matthew Mattina单位Arm ML Research Lab, Boston, MA, USA期刊IEEE Computer Architecture Letters, Vol. 19, No. 1, January-June 2020DOIhttps://doi.org/10.1109/LCA.2020.2979965