Arm Scalable Matrix Extension介绍

最近本人写了一篇介绍Arm Scalable Matrix Extension (可伸缩矩阵扩展，SME)的文章，https://community.arm.com/arm-community-blogs/b/architectures-and-processors-blog/posts/arm-scalable-matrix-extension-introduction。此为中文版，也加入了部分Introducing Armv9 Scalable Matrix Extension for AI Innovation on the Arm CPU https://newsroom.arm.com/blog/scalable-matrix-extension内容。

Armv9-a构架引入了在arm CPU上加速和保护如大语言（LLM）模型这样的高级generative AI应用的功能。Arm Scalable Matrix Extension (SME)是为了满足当前日益复杂和高能耗的AI和机器学习应用需求，创新性地设计的CPU功能。除了加速现今的AI，SME也提供了在arm构架上处理不断更新的generative AI应用的灵活性。

SME显著提升了在arm CPU上处理现有AI和ML应用的能力，从而让各类AI加持的设备和应用有更快更及时的用户体验。它也可以加速很多利用矩阵运算的应用，例如AR，VR和图像处理，它们在AI和ML扮演着越来越重要的角色。
与arm CPU可以处理各类神经网络中诸多数据类型类似，SME提供了不断进化和日益复杂的AI和ML应用需要的灵活性。这保证了arm构架在AI快速演化的年代，仍然承担AI计算重要职责。

Arm SME是一个增强矩阵操作的构架扩展。SME建立在SVE2的基础之上，新增了高效处理矩阵的能力。其关键功能包括：

• 计算两个SVE向量的Outer product (外积）
• 矩阵块（tile）的存储
• 存取矩阵tile中的向量，向矩阵块中插入向量和提取矩阵块里的向量，包括on-the-fly矩阵转置
• Streaming SVE模式（Streaming SVE mode)

下表总结了SME，SVE和SVE2的主要功能：

SME主要支持了以下新功能：

• 新的构架状态（architetecture state),可以用来存储二维（2D）矩阵块。
• 新的Streaming SVE模式(Streaming SVE mode),支持运行SVE2指令，Streaming SVE模式下，SVE2 向量具有矩阵块一样宽度，通常比Non-streaming SVE模式下的向量要更宽。
• 新outer product指令用于计算两个向量的外积，且可以进行外积的累加或累减，将结果放入一个矩阵块。
• 新的存取和移动(load, store, move)指令，用于从矩阵块的一行或是一列读取一个向量到向量寄存器，或是将一个向量写入矩阵块的一行或是一列。

与SVE2一样，SME是一个可伸缩（scalable）向量扩展,它同样支持 Vector Length Agnostic (VLA), per-lane predication, predicate-driven loop control and management这些功能。

Streaming SVE模式

为了支持SME，构架新增了一个操作模式，Streaming SVE mode。在此模式下，包含了现有SVE2指令的一个子集（支持大多SVE2指令），增加了新的SME特有的指令。

Streaming SVE模式支持对大数据集进行高吞吐“streaming”数据处理，被stream的数据通常有较简单的循环控制和有限的条件处理。

在Non-streaming SVE 模式下，支持全部的SVE2指令集，处理一般代码，通常处理复杂的数据结构和复杂的判断。

大多数新加入的SME指令只能在Streaming SVE模式下使用。Streaming vector length (Streaming SVE模式下的硬件向量长度，SVL)可以和Non-streaming SVE mode的向量长度(NSVL)不一样。

预期是：SVL要比NSVL更长或是相同，也就是SVL >= NSVL。例如，NSVL的长度可以为128-bit, 而SVL的长度可以为512-bit。

SME的SVL可以是128-bit, 256-bit,512-bit, 1024-bit或是2048-bit。 SVL需要是2的次幂，而NSVL需要是128的整数倍。

与SVE2类似，软件可以控制SMCR_ELx.LEN寄存器位来设置EL1, EL2, EL3想用的生效SVL长度（可以设置为比硬件支持的SVL更短）。

可以参看Arm Architecture Reference Manual for A-profile architecture构架文档B.1.4.6章节，获取更多Streaming mode的信息。

切换Non-streaming和Streaming SVE模式

如果CPU硬件实现既支持Streaming SVE mode的SME，又支持Non-streaming SVE mode的SVE2，应用程序可以根据自己的需求动态切换这两个操作模式。

让SME有一个独立的操作模式, 使CPU硬件实现可以为同一应用提供不同的向量长度。比如CPU硬件实现可以选择支持一个更长的Streaming SVE mode向量长度，它可以为streaming做硬件优化，适用于高吞吐量的数据处理。

在同一应用中很容易动态切换Streaming和Non-streaming SVE mode。 SME引入了一些新的PSTATE位, PSTATE.{SM, ZA},用于enable/disable Streaming SVE mode和SME ZA storage(有关ZA storage，后面会介绍):

• SM: Enable/disable Streaming SVE mode
• ZA: Enable/disable ZA storage access

可以通过MSR/MRS指令操作Streaming Vector Control Register (SVCR) 来设置和读取PSTATE.{SM, ZA}

• MSR SVCRSM, #<imm>
• MSR SVCRZA, #<imm>
• MSR SVCRSMZA, #<imm>

SMSTART指令是设置PSTATE.SM, PSTATE.ZA的MSR指令的别名。

• SMSTART: Enable both Streaming SVE mode and ZA storage access
• SMSTAT SM : Enable Streaming SVE mode
• SMSTART ZA: Enable ZA storage access

SMSTOP指令是清除PSTATE.SM, PSTATE.ZA的MSR指令的别名。

• SMSTOP: Disable both Streaming SVE mode and ZA storage access
• SMSTOP SM : Disable Streaming SVE mode
• SMSTOP ZA: Disable ZA storage access

下图展示了如何在同一应用在切换Streaming和Non-streaming模式：

可以参看Arm Architecture Reference Manual for A-profile architecture构架文档C6.2.327和C6.2.328章节，获取更多Streaming mode和Non-streaming模式的信息。

SME构架状态

与SVE2类似，在Streaming SVE模式，它有Z0-Z31向量寄存器，和P0-P15 Predicate寄存器。

物理上，SIMD（NEON),Floating point Vn, Qn, Dn, Sn, Hn, Bn寄存器和与之对应的Zn向量寄存器重叠。也就是改写Vn, Qn, Dn, Sn, Hn, Bn寄存器就是改写Zn寄存器的低bits。

构架上要求，当进入Streaming SVE模式时（即PSTATE.SM由0变1）和退出Streaming SVE模式时（即PSTATE.SM由1变0），Zn和Pn寄存器被硬件清零。

大多数原有的SVE2指令还能在Streaming SVE模式下使用，但使用的是SVL长度的向量，与Non-streaming SVE mode向量长度（NSVL）可能不一样。生效的SVL长度可以通过RDSVL指令读出来：

//Read multiple of Streaming SVE vector register size to Xd

RDSVL <Xd>, #<imm>

注意:因为SME支持Vector Length Agnostic (VLA),在Streaming SVE模式下，软件很少需要读SVL向量长度，RDSVL指令通常是在CPU核在Non-streaming SVE模式时，软件读SVL长度。

ZA Array

SME引入了一个新的ZA (Z Array， ZA storage)，它是一个二维（2D）正方形数组，大小是 SVL x SVL。之所以叫Z Array，也是因为它的行和列的长度与Streaming SVE模式的Zn寄存器一致。

例如，如果SVL的长度为256-bit(32-byte)，那么Zn寄存器的长度为32-byte （256-bit)，ZA的大小为 32-byte x 32-byte （或256-bit x 256-bit).

ZA Array可以通过以下方式访问：

• ZA array vectors
• ZA tiles
• ZA tile slices

ZA array vector访问

ZA Array的一行可以当成一个SVL长度的向量(ZA array vector)来访问，这个向量可以放数据类型长度位8-bit, 16-bit, 32-bit, 64-bit, 128-bit的元素：

ZA.B[N], ZA.H[N], ZA.S[N], ZA.D[N], ZA.Q[N]

其中B，H，S，D，Q分别表示8-bit, 16-bit, 32-bit, 64-bit, 128-bit。

ZA array vector的数量与SVL中的byte数相同，例如，如果SLV是256-bit(32-byte),那么ZA array vector的数量就是32个 (0<= N <32)。

为了支持上下文切换，SME引入了新的LDR/STR指令用于从内存中存取一个ZA array vector:

• LDR ZA[<Wv>, <imm>], [<Xn|SP>{, #<imm>, MUL VL}]
• STR ZA[<Wv>, <imm>], [<Xn|SP>{, #<imm>, MUL VL}]

ZA tiles

SME支持对ZA进行分块，分成多个ZA tile进行访问/使用。一个ZA tile是在ZA中的正方形的二维子矩阵。一个ZA tile的宽度总是SVL，这与ZA array的宽度一致。

ZA可以分成多少个ZA tile使用是由元素的数据类型大小决定的：

• 当元素数据类型为8-bit时，ZA只能是一个ZA tile, ZA0.B
• 当元素数据类型为16-bit时，ZA可分为2个ZA tile, ZA0.H-ZA1.H
• 当元素数据类型为32-bit时，ZA可分为4个ZA tile, ZA0.S-ZA3.S
• 当元素数据类型为64-bit时，ZA可分为8个ZA tile, ZA0.D-ZA7.D
• 当元素数据类型为128-bit时，ZA可分为16个ZA tile, ZA0.Q-ZA15.Q

这样做可以充分利用ZA storage, 例如，如果元素数据类型为32-bit浮点数且SVL为256-bit(SVL长度可以放8个32-bit浮点数)，那么ZA的大小为32 x 32 bytes(可以放32行x8列的二维浮点数数值), 一个Z寄存器可以放8个浮点数元素的向量。那么两个Z寄存器里的向量做外积运算，产生的外积结果是8行x 8列的二维浮点数数组，这个外积只需要1/4的ZA storage。为了充分利用,这样ZA可以分成4个ZA tile, ZA0.S-Z3.S来使用。

再例如，如果SVL为256-bit(32 byte)，元素的数据类型大小为8-bit,那么ZA可以看成为一个有32行，每行为32个8-bit的数组。

如果SVL同为256-bit(32 byte)，但元素的数据类型大小为16-bit，那么ZA可以看出是两个ZA tile, ZA0.H和ZA1.H, 每个tile由16行，每行有16个16-bit的数组组成。

ZA tile的访问

一个ZA tile可以作为一个整体来访问，也可以以一个个ZA tile slice的方式访问。

当作为一个整体访问时，指令可以使用tile的名字访问：

ZA0.B, ZA0.H-ZA1.H, ZA0.S-ZA3.S, ZA0.D-ZA7.D or ZA0.Q-ZA15.Q

一个ZA tile slice是由其ZA tile中水平方向或是垂直方向的连续元素组成的一维数组，即在ZA tile中的一行或是一列。

对一个ZA tile的向量访问即是读写一个ZA tile slice,

• 是水平或是垂直方向的ZA tile slice访问，由加在ZA tile名字后的'H'或是'V'后缀来表示
• 是哪一个ZA tile slice由ZA tile名字后面的slice索引'[N]'来表示

例如，如果SVL为128-bit且元素类型的大小为8-bit,那么其水平和垂直的ZA tile slices可以由下图展示，

ZA0V.B[0]和ZA0V.B[13]访问ZA0.B ZA tile垂直方向向量(即一列），ZA0H.B[0]和ZA0H.B[15]访问ZA0.B ZA tile水平方向向量(即一行)。

如果SVL为128-bit且元素类型的大小为16-bit,那么其水平和垂直的ZA tile slices可以由下图展示，

事实上，为了高效的硬件对ZA tile和tile slice的访问，ZA tile的tile slice在ZA中以交织方式存储的。

例如，如果SVL是256-bit, 元素数据类型大小是16-bit,那么其水平ZA tile slice在ZA中的存储可以展示为：

如果SVL同样是256-bit, 元素数据类型大小是32-bit,那么其水平ZA tile slice在ZA中的存储可以展示为：

下图展示了不同的元素数据类型大小的水平和垂直方向ZA tile slice的混合视图:

图中左边的那些列显示了ZA storage的每行可以被各种数据类型大小的tile水平slice访问的方式。
设‘SIZE’为向量元素的类型大小，B，H，S，D和Q类型的SIZE分别为1，2，4，8，16。
设'NUM_OF_ELEMENTS'为可以放在向量里的元素个数，即(SVL的byte数）/SIZE。那么，

水平tile slice，ZAnH.<B|H|S|D|Q>[m],访问ZA storage的第(mSIZE+n)行（完整一行）。例如，ZA2H.S[5]访问ZA storage的第（54+2）行，即第22行。ZA6H.D[0]访问ZA storage的第（0*8+6）行，即第6行。

垂直tile slice，ZAnV.<B|H|S|D|Q>[m],访问ZA storage中（i*SIZE+n）这些行里第m个元素（元素大小为SIZE），由这些元素组成的向量，其中 i为 0到(NUM_OF_ELEMNTS-1)。例如，ZA3V.S[4]访问ZA storage中，由第3，7，11，15，19，23，27，31这些行中的第4个元素（元素大小为4）组成的向量。

如果软件中使用混合数据类型大小或是水平/垂直方向tile slice时，需要小心处理tile slice间的重叠。

可以参看Arm Architecture Reference Manual for A-profile architecture构架文档B1.4.8到B1.4.12章节，获取更多ZA，ZA tile, ZA tile slice的信息。

Steaming SVE模式下支持的指令

有些NEON/SVE2指令在Streaming SVE模式受到影响；

• 有些SVE/SVE2指令变为非法执行

  • Gather/Scatter load/store SVE2指令
  • 使用First Fault 寄存器的SVE2指令
• 大多的NEON指令变为UNDEFINED

如果想得到更多所影响指令的信息，请参考 Arm Architecture Reference Manual for A-profile architecture构架文档。

SME主要增加以下指令：

• 矩阵外积并且累加或累减指令，包括 FMOPA, UMOPA, BFMOPA。这些指令利用，

  • SVE2向量寄存器(Z0-Z31)作为外积运算的行和列输入
  • ZA storage用来保存二维的矩阵块输出
• 将SVE2 Z向量与ZA的行或列做加法运算的指令
• 对ZA tiles的清零操作指令
• 增加了一些在Streaming和Non-streaming模式下都能使用的指令

我将在下一篇博客中介绍这些指令。

进一步阅读

有关SME的更具体信息，请参照Arm Architecture Reference Manual for A-profile architecture构架文档：https://developer.arm.com/documentation/ddi0487 。
关于如何在你的应用中使用SME来高效处理矩阵和其他类型的数据，请参照SME Programmer's Guide：https://developer.arm.com/documentation/109246 。