浅析Armv9-A构架上的可伸缩矩阵扩展(SME)

Arm构架为从超级计算机到广泛的设备带来可伸缩的向量处理，使世界上大多数的计算负载都可以运行在arm构架上。 在Arm Vision Day上，arm公布了Armv9-A 构架，随后arm公开了一个Armv9-A构架扩展的早期技术细节，这个扩展就是可伸缩矩阵扩展（Scalable Matrix Extension, SME)。SME是一序列增强CPU构架对矩阵运算支持的最新更新。

这次公开SME技术的目的是提供信息和赋能OS和工具开发者生态。SME为支持矩阵运算引入了一个新的编程模式和寄存器状态。

SME是在可伸缩向量扩展（Scalable Vector Extensions， SVE和SVE2）的基础上建立的，并增加了有效处理矩阵的能力，主要功能包括：

矩阵tile的存储

存取，插入，提取 tile向量，包括on-the-fly 转置

计算SVE向量的外积（Outer product）

Streaming SVE 模式

Arm上的矩阵乘

矩阵乘是很多关键负载的重要部分，比如科学计算，计算机视觉，有些机器学习和增强现实场景。Arm构架已经持续演进，增加了提高这些运算的性能和效率的指令。

图片1

Armv8.4-A:支持了8-bit整型的点积运算指令（DOT指令）
Armv8.6-A: 支持在向量寄存器里做整型和浮点型的矩阵乘指令，和BFloat16数据类型
Armv9-A: 支持了SVE2，更宽的向量处理
SME是这个演化的下一步，可以明显提升CPU处理矩阵的吞吐量和效率。

SME和矩阵乘

一个简单实现矩阵乘的方式是如下图所示的三层嵌套循环，

图片2

这种方式得到的乘运算和内存读的数量比率是1：2，也就是每两个元素读操作对应一个乘运算。为提高效率和吞吐量，我们需要提高这个比率。一个方法就是一次计算多个结果，比如：

在这个例子里面，通过一次计算4个结果可以提高乘运算和读操作的比率到1：1，也就是，得到4个乘结果需要四个读操作。

图片3

SME是基于外积计算引擎，进一步利用一次读操作产生多个结果的想法，

图片4

计算一次向量外积A[H] ⨷ B[W]， 会产生一个HxW的矩阵， 并累加到一个矩阵tile - C[HxW]. 通过迭代计算A矩阵的列和B矩阵的行的外积，并累加到C矩阵，可以得到一个完整的A[HxK] 和B[KxW]矩阵乘结果。

SME是可伸缩的构架设计，允许不同的向量宽度的实现。乘运算和读操作的比率取决于宽度的实现。比如，512-bit 向量的实现，如果向量中存放32-bit的元素，那么其乘操作与读操作的比率为256:2 。当4个向量输入可以计算出4个输出tile时，这个比率提升到256:1 。

SME和SVE2

图片5

SME增加了一新的操作模式 - Streaming SVE 模式。当在Streaming SVE 模式时，可以使用新的SME存储和指令， 同时也可以使用大部分的现有SVE2指令。当在非Streaming SVE模式时，其行为与SVE2完全一样。应用可以在其需要时切换操作模式。

为SME操作设计单独的操作的模式使得硬件实现为Streaming和非Streaming模式设计不同向量长度，在一个应用中可以使用Streaming和非Streaming 模式。 比如一个实现可以为Streaming SVE模式设计一个更长的向量，硬件优化Streaming, 以吞吐量为中心的操作。

原作者：Martin Weidmann