此文为Part 2: Arm Scalable Matrix Extension (SME) Instructions 的中文版。
前文简介了SME，本文介绍SME指令。
Arm Scalable Matrix Extension介绍
操作ZA storage的SME指令主要包括：
- 计算两个向量外积，并累加或累减，将结果放入一个ZA tile的指令
- 内存与ZA tile行或列之间的存取操作指令，以及SVE Z 寄存器和ZA tile 行或列之间的移动指令
- 水平或垂直方向上，一个Z向量与ZA tile的加法指令
- 将一个标量寄存器加上Streaming SVE模式的向量长度倍数的指令
外积并累加或累减指令
为了帮助理解外积并累加或累减指令，让我们看看如何使用外积操作来做矩阵乘法。下图展示了外积运算：
计算a和b两个向量的外积，可以得到他们外积的结果-矩阵C。
现在我们看看如果进行矩阵a和b的矩阵乘：
这个矩阵乘可以通过计算两次外积操作再累加两个外积结果来实现，如下图所示：
SME为以下数据类型引入了高效的外积并累加或累减指令：
• 8-bit, 16-bit整数
• FP16,BF16, FP32和FP64浮点数
这些指令计算放在两个Z向量寄存器（Zn和Zm）里的向量的外积，并将这个外积结果累加或累减一个ZA tile（ZAda）中已有数据，将结果存入同一ZA tile (ZAda)。每个源向量都可以被其相应的控制predicate寄存器（Pn和Pm）独立地predicate（有关predicate的含义，请阅读SVE相关文章）。
FP32, FP64外积并累加或累减指令
那些输入向量和输出数组有同样数据类型（FP32， FP64）的指令相对直观。
下例展示了FP32类型的外积并累加或累减指令。
FMOPA <ZAda>.S, <Pn>/M, <Pm>/M, <Zn>.S, <Zm>.S
FMOPS <ZAda>.S, <Pn>/M, <Pm>/M, <Zn>.S, <Zm>.S
这个例子中，假设SVL向量长度为128，Zn.S和Zm.S中存放了4个FP32数组成的向量，此指令计算Zn.S和Zm.S的外积，外积结果为图中灰色的矩阵，然后将此外积结果累加或累减ZAda.S这个ZA tile（矩阵）中原有的值，将结果存入同一ZAda.S tile。
FP16, BF16, INT16, INT8, I16I64类型的外积并累加或累减指令
为了保持这些数据类型结果的精度，以及充分利用计算硬件资源和ZA storge，这些指令会扩大计算结果数据类型，因此这些操作不像前面FP32，FP64类型指令那么直接。
- BF16指令计算两个BF16的外积的和，扩大结果类型为FP32, 然后这些结构被解构性地从目标tile中加或减。
- INT8指令计算四个INT8的外积的和，扩大结果类型为INT32, 然后这些结构被解构性地从目标tile中加或减。
- INT16指令计算两个INT16的外积的和，扩大结果类型为INT32, 然后这些结构被解构性地从目标tile中加或减。
- FP16指令计算两个FP16的外积的和，扩大结果类型为INT32, 然后这些结构被解构性地从目标tile中加或减。
- 如果实现了FEAT_SME_I16I64，I16I64指令计算四个INT16的外积的和，扩大结果类型为INT64, 然后这些结构被解构性地从目标tile中加或减。
以下例子展示了SVL向量长度为128的INT8 UMOPA指令进行的操作：
UMOPA <ZAda>.S, <Pn>/M, <Pm>/M, <Zn>.B, <Zm>.B
对于UMPOA指令，每个输入向量(Zn.B, Zm.B)被当成一个有4x4元素的矩阵，可以看作是4个连续的元素组成的块（如图中红线标出的那样）被转置了。
在这个例子中， 因为SVL向量长度为128-bit：
- 第一源向量，Zn.B，包含一个无符号8-bit整数的4x4子矩阵。
- 第二源向量，Zm.B，包含一个无符号8-bit整数的4x4子矩阵。
- UMOPA指令计算出4x4扩大了的32-bit整数外积的和，然后解构性地加上目标tile，ZAda中的整数。
更通用地说，这条无符号整型外积和并累加指令（UMOPA）将第一源向量中的子矩阵乘于第二源矩阵中的子矩阵。每个源向量包含一个(SVL/32) x 4的无符号8-bit整数的子矩阵。然后将得到的 (SVL/32) x (SVL/32)扩大了的32-bit整数外积和解构性地加上一个32-bit整数目标tile。
下面的例子展示了SVL为128-bit的BF16 BFMOPA进行的操作：
BFMOPA <ZAda>.S, <Pn>/M, <Pm>/M, <Zn>.H, <Zm>.H
这个例子中，因为SVL为128-bit:
- 第一源向量，Zn.H，包含一个BF16浮点数的4x2子矩阵，它被扩大为单精度浮点数。
- 第二源向量，Zm.H，包含一个BF16浮点数的2x4子矩阵，它被扩大为单精度浮点数。
- BFMOPA指令计算出4x4 单精度外积的和，然后解构性地加上目标tile，ZAda中的单精度数。
更通用地说，这条指令（BFMOPA）扩大了存放在第一源里的(SVL/32) x2 BF16子矩阵的类型为单精度，扩大了存放在第二源里的2x (SVL/32) BF16子矩阵的类型为单精度，将这两个子矩阵相乘。然后将得到的 (SVL/32) x (SVL/32)单精度外积和解构性地加上一个单精度目标tile。
以下表格显示了几种数据类型和SVL长度的一条外积并累加或累减指令所做的对应数据类型的MAC(乘累加)数量：
带Predication的SME指令
每个源向量都可以被其相应的控制predicate寄存器独立地predicate:
- 外积并累加或累减指令使用Pn/M和Pn/M (没有/Z形式)：Inactive的源元素被当成具有0值。
- Slice move指令使用Pg/M: 目标slice中Inactive的元素保持不变。
- Tile slice load指令使用Pg/Z: 目标tile slice中的Inactive元素被设置为0。
- Tile slice store指令使用Pg: Inactive的元素不会写入内存。
Predication让矩阵的维数不是SVL的倍数的情况更容易处理。
如下例所示：
输入向量Z0被P0 predicate, Z1被P1 predicate.
在这个例子里：
- SVL为512-bit
- Z寄存器包含16个FP32数组成的向量
- P0中最后两个元素是Inactive的
- P1中最后一个元素是Inactive的
这条指令更新ZA0.S中(16-2) x (16-1) 个FP32元素，因为使用了Pn/M , ZA0.S中剩下的元素保持不变。
下图展示了更多的predicated外积并累加或累减的例子。图中被划线的文字表示被inactive predicate元素影响的计算部分。
ZA tile与一个Z向量的加运算
SME包括ZA tile的行或列都加上一个向量的指令，这些指令也有predication的支持。
例如：
ADDHA ZA0.S, P0/M, P1/M, Z1.S 它进行以下操作：
这个ADDHA指令将源向量Z1中的每个元素加上ZA0.S tile每一水平slice的相应active元素。
Tile中元素被一对governing predicate进行predicate。 一个水平slice中的一个元素在下面情况下可以认为是active：
- 它在第二governing predicate对应的元素是TRUE, 并且
- 它在第一governing predicate对应的水平slice行号也为TRUE,目标tile中inactive元素保持不变。
Tile load, store, move指令
SME tile load, store, move指令可以：
- 从内存读取数据，放入ZA tile的行或列
- 将ZA tile的行或列写入内存
- 将ZA tile的行移动到SVE Z向量寄存器
- 将SVE Z向量寄存器移动到ZA tile行或列
Tile slice load和store指令
LD1B, LD1H, LD1S, LD1D和 LD1Q指令从内存中连续读取一些值到ZA tile slice，数据类型分别是8-bit, 16-bit, 32-bit, 64-bit和128-bit元素。
ST1B, ST1H, ST1S, ST1D和ST1Q指令将ZA tile slice存入连续内存中，数据类型分别是8-bit, 16-bit, 32-bit, 64-bit和128-bit元素。
这些指令也用predication的支持，例如：
LD1B ZA0H.B[W0, #imm], P0/Z, [X1, X2] 此LD1B指令执行predicated的连续byte读取，它从地址为(X1+X2)的内存读取数据到ZA0中行号为（W0+imm）的这个水平tile slice中。目标tile slice中Inactive的元素被设置为0.
ST1H ZA1V.H[W0, #imm], P2, [X1, X2, LSL #1]此ST1H指令执行predicated连续halfword的存操作，它将ZA1中列号为（W0+imm）的垂直 tile slice存到地址为（X1+X2*2）的内存， tile slice中Inactive的元素不写入内存。
Tile slice move指令
MOV指令（MOVA指令的别名）将一个Z向量寄存器的值移动到一个ZA tile slice，或将一个ZA tile slice中的值移动到一个Z向量寄存器。这条指令操作带指定元素大小的ZA tile的单个水平或垂直tile slice。 Slice的行号/列号由slice的检索寄存器加上立即数偏移指定。目标slice中Inactive的元素保持不变。
例如：
MOV ZA0H.B[W0, #imm], P0/M, Z0.B或
MOVA ZA0H.B[W0, #imm], P0/M, Z0.B此指令将向量寄存器Z0.B中的值移动到ZA0H.B[W0,#imm]这个水平ZA tile slice中，使用P0作为predication寄存器。目标tile slice中Inactive的元素保持不变。
ZA array vector load/store指令
SME LDR指令从内存读取数据到一个ZA array向量，SME STR指令将一个ZA array向量中的值存入内存。
这些指令是不带predication功能的。它们主要是为了软件的context switching时对ZA storage进行save/restore。SME LDR/STR指令也可以在Non-streaming SVE模式下，当PSTATE.ZA使能的情况下使用。
例如，下面的STR指令的ZA array向量是由一个向量选择寄存器Wv（标量寄存器W）加上可选的立即数（Wv+Imm）指定。访问内存的地址为：一个标量寄存器作为base，加上相同的可选立即数偏移乘以当前向量长度byte数。
STR ZA[<Wv>, <imm>], [<Xn|SP>{, #<imm>, MUL VL}]ZA tile清零指令
SME ZERO指令可以清零一组64-bit ZA tile:
ZERO { <mask>}ZERO指令可以清零多到8个名为ZA0.D到ZA8.D的ZA tile，那些tile要清零由指令中的mask指定，剩下的其他tile保持不变。
这条指令也可以在Non-streaming SVE模式，当PSTATE.ZA使能的情况下使用。
如果要清零整个ZA array, 可以使用一个指令别名，ZERO {ZA}。
新的SVE2指令
SME构架扩展加入了一些新的SVE2指令，这些指令也可以在PE实现了SVE2, 处于Non-streaming SVE模式时使用。这些指令包括：
- 选择一个predicate寄存器或是all-false的Predicate select指令
- 翻转（Reverse）64-bit double word元素的指令
- 有符号/无符号钳位为更小/更大值向量的指令
下面介绍以下Predicate select指令。
PSEL指令
PSEL指令选择一个predicate寄存器或是all-false到目标predicate寄存器，如下所示：
PSEL <Pd>, <Pn>, <Pm>.<T>[<Wv>, <imm>]如果指令中第二源predicate寄存器（Pm）中指定的元素为True, 这条指令将第一源predicate寄存器(Pn)的内容放到目标predicate寄存器(Pd), 否者设置目标predicate寄存器的值全部为false。
例如以下指令，假设W12的值为0：
PSEL P0, P1, P2.B[W12, #0]第二源predicate寄存器的[W12+0]即[0]个元素为False, 因此目标寄存器P0被设置为全0（all-false），如下图所示：
现在看看如下指令，仍然假设W12的值为0，但这次立即数偏移为1：
PSEL P0, P1, P2.B[W12, #1]第二源predicate寄存器的[W12+1]即[1]个元素为True, 因此选择第一源predicate寄存器的值到目标寄存器P0，如下图所示：
进一步阅读
有关SME更详细信息，请参看 Arm Architecture Reference Manual for A-profile architecture.
关于如何在你的应用中高效使用SME来处理矩阵和其他类型数据，请参看SME Programmer's Guide.
