在高性能计算和深度学习领域，通用矩阵乘法（GEMM）作为最核心的计算内核之一，其性能优化始终是业界关注的焦点。随着 ARM 架构在数据中心和边缘计算场景的广泛应用，特别是 ARM Scalable Matrix Extension（SME）扩展的推出，为矩阵运算带来了硬件级别的加速能力。然而，现有线性代数库在 SME 上的实现往往未能充分释放其性能潜力，尤其是在处理大规模矩阵时表现欠佳。本文将深入探讨 ARM SME 的矩阵分块与向量外积指令特性，并详细阐述面向 GEMM 的寄存器分配与流水线调度优化策略。
ARM SME 架构建立在可扩展向量扩展（SVE）之上，引入了一套专为矩阵运算优化的指令集和硬件单元。其核心特性包括流式 SVE 模式（SSVE）、二维 ZA 存储区以及外积处理指令。ZA 存储区是一个大小为 SVL×SVL 的二维矩阵寄存器，在 Apple M4 Pro 平台上，SVL 固定为 512 位，使得每个 Z 寄存器大小为 64 字节，整个 ZA 存储区达到 64×64=4096 字节。对于 32 位元素，ZA 存储可视为四个独立 tile（ZA0.S 至 ZA3.S），每个 tile 大小为 16×16；而对于 64 位元素，则可划分为八个 tile（ZA0.D 至 ZA7.D），每个 tile 大小为 8×8。这种灵活的 tile 划分方式为矩阵运算提供了丰富的并行度。
SME 的外积累加指令是实现高效 GEMM 的关键。以 FP32 精度为例，FMOPA 指令计算两个长度为 SVL/32 的向量的外积，并将结果累加到 ZA tile 中，每次操作可执行 2×(SVL/32)^2 次浮点运算。更重要的是，SME 支持低精度输入高精度输出的混合精度操作，如 FP16 或 BFloat16→FP32、INT8→INT32、INT16→INT32/INT64。这种特性显著提升了计算强度，减少了内存带宽压力，使得混合精度 GEMM 成为可能。
然而，现有开源 SME 优化库在实现中存在三个主要瓶颈。首先，缓存层次结构利用不足：这些库普遍采用简单的三级嵌套循环，tile 大小设计未能有效匹配共享 L2 缓存，导致数据局部性差。其次，数据打包方法次优：LIBXSMM 和 OpenBLAS 仅打包一个输入矩阵以减少打包开销，但当未打包矩阵超过共享 L2 缓存容量时，会引发大量缓存未命中和 TLB 缺失。最后，多向量加载指令利用不充分：SME2 支持多向量数据处理，但现有微内核仅使用单寄存器或双寄存器加载，带宽效率低下。
为解决这些瓶颈，我们提出了一套系统化的优化策略。在缓存感知分区方面，我们引入分析模型指导大规模矩阵的划分，确保子矩阵尺寸适配共享 L2 缓存。通过将传统的三级循环结构转换为六级阻塞算法，优化缓存局部性和 TLB 效率。设 L2 缓存大小为 C，数据类型大小为 d，则阻塞参数需满足约束：mc×kc + kc×nc + mc×nc < C/d。同时，通过拉格朗日乘数法最大化计算与内存比（CMR）来优化阻塞参数。
在数据打包策略上，我们采用双矩阵打包和即时转置技术。对于矩阵 A，利用 ZA tile 的水平和垂直切片能力实现即时转置：将行向量加载到水平切片，再通过垂直切片写出，实现从行主序到列主序的高效转换。对于矩阵 B，采用首轮在线打包策略：在 FMOPA 指令执行的同时，并行进行数据打包，将 kc×nr 面板顺序写入线性缓冲区，有效隐藏内存访问延迟。
在微内核设计层面，我们专注于最大化向量加载粒度和 tile 利用率。微基准测试表明，四寄存器加载（256 字节）可实现高达 900 GB/s 的带宽，而单寄存器加载（64 字节）仅为 230 GB/s。因此，我们设计主微内核采用 mr×nr=16×64 的 tile 尺寸（FP32 精度下），每次加载四个 Z 寄存器，确保同时更新所有四个 ZA tile。微内核执行流程如下：首先将 Cr 子块以四寄存器方式加载到 ZA tile；然后分批次加载 Ar 和 Br 数据，每个批次包含四个 Z 寄存器；在循环内通过循环展开和软件流水线技术，交替执行加载和 FMOPA 操作，隐藏内存访问延迟；最后将结果写回 Cr 子块。
对于边界处理，我们设计了专用的边缘微内核。当矩阵维度不是微内核尺寸的整数倍时，使用谓词寄存器遮蔽尾部元素，避免性能损失。混合精度 GEMM（如 FP16 输入→FP32 输出）进一步利用了 SME 的原生支持，通过将两个 FP16 元素视为一个 FP32 值进行打包和计算，在保持精度的同时提升吞吐量。
在 Apple M4 Pro 平台上的实验评估表明，优化后的实现相比 Apple Accelerate 库实现了平均 1.23 倍的加速，相比 LIBXSMM、KleidiAI 和 OpenBLAS 分别达到 3.96 倍、4.69 倍和 5.7 倍的性能提升。混合精度 FP16 GEMM 可达到 SME 峰值性能的 94%，INT8 GEMM 更是实现了高达 7856 GOPS 的吞吐量。这些结果充分验证了寄存器分配优化和流水线调度策略的有效性。
总结而言，深入理解 ARM SME 架构特性，针对性地设计矩阵分块、数据打包和微内核实现，是释放其 GEMM 性能潜力的关键。通过采用缓存感知分区、双矩阵即时打包和四向量加载微内核等策略，可以显著提升密集计算性能，为 ARM 架构在高性能计算和深度学习领域的应用奠定坚实基础。
资料来源：本文主要参考《Demystifying ARM SME to Optimize General Matrix Multiplications》（arXiv:2512.21473）及相关 ARM SME 技术文档。