Arm 可伸缩矩阵扩展 (SME) 作为 Armv9 架构中的一项创新特性，旨在满足当前日益复杂和高能耗的人工智能 (AI) 和机器视觉 (ML) 应用需求。除了加速现今的 AI，SME 也提供了在 Arm 架构上处理不断更新的生成式 AI 应用的灵活性。在上一篇内容中，Arm 技术专家为大家简要介绍了 SME，本周我们将带各位更为详细地来了解 SME 的指令，助力你在应用中高效使用 SME！
操作 ZA 存储的 SME 指令主要包括：计算两个向量外积，并累加或累减，将结果放入一个 ZA tile 的指令；内存与 ZA tile 行或列之间的存取操作指令，以及 SVE Z 寄存器和 ZA tile 行或列之间的移动指令；水平或垂直方向上，一个 Z 向量与 ZA tile 的加法指令；将一个标量寄存器加上 Streaming SVE 模式的向量长度倍数的指令。
外积并累加或累减指令
为了帮助理解外积并累加或累减指令，让我们看看如何使用外积操作来做矩阵乘法。下图展示了外积运算：
计算 a 和 b 两个向量的外积，可以得到他们外积的结果，即矩阵 C。
现在我们看看如果进行矩阵 a 和 b 的矩阵乘：
这个矩阵乘可以通过计算两次外积操作再累加两个外积结果来实现，如下图所示：
SME 为 8 位、16 位整数，以及 FP16、BF16、FP32 和 FP64 浮点数引入了高效的外积并累加或累减指令。这些指令计算放在两个 Z 向量寄存器（Zn 和 Zm）里的向量的外积，并将这个外积结果累加或累减一个 ZA tile (ZAda) 中已有数据，将结果存入同一 ZA tile (ZAda)。每个源向量都可以被其相应的控制 predicate 寄存器（Pn 和 Pm）独立地 predicate。
FP32、FP64 外积并累加或累减指令
那些输入向量和输出数组有同样数据类型（FP32 和 FP64）的指令相对直观。下面这一例子展示了 FP32 类型的外积并累加或累减指令。
这个例子中，假设 SVL 向量长度为 128 位，Zn.S 和 Zm.S 中存放了四个 FP32 数组成的向量，此指令计算 Zn.S 和 Zm.S 的外积，外积结果为图中灰色的矩阵，然后将此外积结果累加或累减 ZAda.S 这个 ZA tile（矩阵）中原有的值，将结果存入同一 ZAda.S tile 中。
FP16、BF16、INT16、INT8、 I16I64 类型的外积并累加或累减指令
为了保持这些数据类型结果的精度，以及充分利用计算硬件资源和 ZA 存储，这些指令会扩大计算结果数据类型，因此这些操作不像前面 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 指令进行的操作：
对于 UMPOA 指令，每个输入向量 (Zn.B, Zm.B) 被当成一个有 4x4 元素的矩阵，可以看作是四个连续的元素组成的块（如图中红线标出）被转置了。
在这个例子中，因为 SVL 向量长度为 128 位，第一源向量 Zn.B 包含一个无符号 8 位整数的 4x4 子矩阵，第二源向量 Zm.B 包含一个无符号 8 位整数的 4x4 子矩阵，UMOPA 指令计算出 4x4 扩大了的 32 位整数外积的和，然后解构性地加上目标 tile (ZAda) 中的整数。
更通用地说，这条无符号整型外积和并累加指令 (UMOPA) 将第一源向量中的子矩阵乘以第二源向量中的子矩阵。每个源向量包含一个 (SVL/32) x 4 的无符号 8 位整数的子矩阵。然后将得到的 (SVL/32) x (SVL/32) 扩大了的 32 位整数外积和解构性地加上一个 32 位整数目标 tile。
下面的例子展示了 SVL 为 128 位的 BF16 BFMOPA 进行的操作：
这个例子中，因为 SVL 为 128 位，第一源向量 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 位，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 的行或列都加上一个向量的指令
它进行以下操作：
这个 ADDHA 指令将源向量 Z1 中的每个元素加上 ZA0.S tile 每一水平 slice 的相应 active 元素。Tile 中元素被一对控制 predicate 进行 predicate。一个水平 slice 中的一个元素在下面情况下可以认为是 active：它在第二控制 predicate 对应的元素是 true，并且它在第一 控制 predicate 对应的水平 slice 行号也为 true，目标 tile 中 inactive 元素保持不变。
Tile 存取和移动指令
SME tile 存取和移动 (load, store, move) 指令可以：
- 从内存读取数据，放入 ZA tile 的行或列
- 将 ZA tile 的行或列写入内存
- 将 ZA tile 的行移动到 SVE Z 向量寄存器
- 将 SVE Z 向量寄存器移动到 ZA tile 行或列
Tile slice 存取指令
LD1B、LD1H、LD1S、LD1D 和 LD1Q 指令从内存中连续读取一些值到 ZA tile slice，数据类型分别是 8 位、16 位、32 位、64 位和 128 位元素。ST1B、ST1H、ST1S、ST1D 和 ST1Q 指令将 ZA tile slice 存入连续内存中，数据类型分别是 8 位、16 位、32 位、64 位和 128 位元素。这些指令也用 predication 的支持，例如：
此 LD1B 指令执行 predicated 的连续字节读取，它从地址为 (X1+X2) 的内存读取数据到 ZA0 中行号为 (W0+imm) 的这个水平 tile slice 中。目标 tile slice 中 inactive 的元素被设置为 0。
此 ST1H 指令执行 predicated 连续半字的存操作，它将 ZA1 中列号为 (W0+imm) 的垂直 tile slice 存到地址为 (X1+X2*2) 的内存，tile slice 中 inactive 的元素不写入内存。
Tile slice 移动指令
MOV 指令（MOVA 指令的别名）将一个 Z 向量寄存器的值移动到一个 ZA tile slice，或将一个 ZA tile slice 中的值移动到一个 Z 向量寄存器。这条指令操作带指定元素大小的 ZA tile 的单个水平或垂直 tile slice。Slice 的行号/列号由 slice 的检索寄存器加上立即数偏移指定。目标 slice 中 inactive 的元素保持不变。
此指令将向量寄存器 Z0.B 中的值移动到 ZA0H.B[W0, #imm] 这个水平 ZA tile slice 中，使用 P0 作为 predication 寄存器。目标 tile slice 中 inactive 的元素保持不变。
ZA array 向量存取指令
SME LDR 指令从内存读取数据到一个 ZA array 向量，SME STR 指令将一个 ZA array 向量中的值存入内存。这些指令是不带 predication 功能的。它们主要是为了软件上下文切换时，对 ZA 存储进行保存或恢复。SME LDR/STR 指令也可以在 Non-streaming SVE 模式下，当 PSTATE.ZA 使能的情况下使用。例如，下面的 STR 指令的 ZA array 向量是由一个向量选择寄存器 Wv（标量寄存器 W）加上可选的立即数 (Wv+Imm) 指定。访问内存的地址为：一个标量寄存器作为基础，加上相同的可选立即数偏移乘以当前向量长度字节数。
ZA tile 清零指令
SME 清零 (ZERO) 指令可以清零一组 64 位 ZA tile：
清零指令可以清零多到八个名为 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 选择指令
- 翻转 (reverse) 64 位双字元素的指令
- 有符号/无符号钳位为更小/更大值向量的指令
PSEL 指令
PSEL 指令选择一个 predicate 寄存器或是 all-false 到目标 predicate 寄存器，如下所示：
如果指令中第二源 predicate 寄存器 (Pm) 中指定的元素为 true，这条指令将第一源 predicate 寄存器 (Pn) 的内容放到目标 predicate 寄存器 (Pd)，否则设置目标 predicate 寄存器的值全部为 false。例如以下指令，假设 W12 的值为 0：
第二源 predicate 寄存器的 [W12+0] 即 [0] 个元素为 false，因此目标寄存器 P0 被设置为全 0 (all-false)，如下图所示：
现在来看如下指令，仍然假设 W12 的值为 0，但这次立即数偏移为 1：
第二源 predicate 寄存器的 [W12+1] 即 [1] 个元素为 true，因此选择第一源 predicate 寄存器的值到目标寄存器 P0，如下图所示：
作者：Arm社区
文章来源：Arm社区
推荐阅读
