

Arm 矩阵加速：可扩展矩阵扩展 SME

ai算法芯片与系统

2025-11-13

导读：本文介绍了 Arm 架构中的可扩展矩阵扩展（SME），重点阐述了其在流式 SVE 模式下的高效矩阵计算能力，以及利用 ZA 存储阵列进行大规模数据存储和灵活访问的机制，为高性能计算应用提供了强大的硬件

1. SME 简介

可扩展矩阵扩展 SME 建立在可扩展向量扩展（SVE 和 SVE2）的基础上，增加了高效处理矩阵的能力。其主要特性包括：

• 计算 SVE 向量的外积
• 矩阵分片存储
• 加载、存储、插入和提取分片向量（包括动态转置）
• 流式 SVE 模式

下表总结了 SME、SVE 和 SVE2 的主要特性：

SME	SVE	SVE2
流式 SVE 模式	NEON DSP++	可扩展向量
动态矩阵转置	多精度算术	逐通道谓词
向量外积	匹配检测和直方图	聚集加载和分散存储
加载、存储、插入和提取矩阵向量	非临时性分散/聚集	预测向量化
位操作置换	ML 扩展 (FP16 + DOT)
AE, SHA3, SM4, 加密	V8.6 BF16, FP 和 Int8 支持

SME 定义了以下新特性：

• 新的架构状态，可用于存储二维矩阵分片。
• 流式 SVE 模式，支持执行向量长度与分片长度匹配的 SVE2 指令。
• 新的指令，用于将两个向量的外积累加（或相减）到矩阵分片中。
• 新的加载、存储和移动指令：向量可以被写入矩阵分片的行或列，或者矩阵分片的行或列可以被读入向量。

与 SVE2 类似，SME 也是一个支持可扩展向量长度的扩展，实现了向量长度不可知（VLA）、逐通道谓词、谓词驱动的循环控制和管理功能。

2. 流式 SVE 模式

SME 引入了流式 SVE 模式，它实现了 SVE2 指令集的一个子集，并增加了新的 SME 专用指令。

• 流式 SVE 模式支持针对大型数据集的高吞吐量流式数据处理，流式数据通常具有简单的循环控制流和有限的条件性。
• 在非流式 SVE 模式下，支持完整的 SVE2 指令集，通常处理复杂的数据结构和复杂的判断。

大多数 SME 指令仅在流式 SVE 模式下可用。流式 SVE 模式下的流式向量长度（SVL） 可能与非流式向量长度（NSVL） 不同。

期望是：SVL 应大于或等于 NSVL，即 SVL >= NSVL。例如，NSVL 的长度可以是 128 位，而 SVL 的长度可以是 512 位。

SME 的 SVL 可以是 128 位、256 位、512 位、1024 位或 2048 位。SVL 需要是 2 的幂，NSVL 需要是 128 的倍数。

与 SVE2 类似，软件可以通过控制 SMCR_ELx.LEN 寄存器位来设置 EL1、EL2、EL3 希望使用的有效 SVL 长度（可以设置得比硬件支持的 SVL 更短）。

有关流式 SVE 模式的更多信息，请参阅 Arm 架构参考手册（A-profile 架构）的 B1.4.6 节。

3. 非流式与流式 SVE 模式之间的切换

如果 CPU 硬件实现同时支持 SME 的流式 SVE 模式和 SVE2 的非流式 SVE 模式，应用程序可以根据自身需求在这两种操作模式之间动态切换。

• 为 SME 提供独立的操作模式，允许 CPU 硬件实现为同一应用程序提供不同的向量长度。例如，CPU 硬件实现可以选择支持更长的流式 SVE 模式向量长度，并针对适合高吞吐量的流操作优化硬件。
• 应用程序可以轻松地在流式 SVE 模式和非流式 SVE 模式之间动态切换。SME 引入的 PSTATE.{SM, ZA} 位可以启用和禁用流式 SVE 模式和 SME ZA 存储：

• SM：启用和禁用流式 SVE 模式
• ZA：启用和禁用 ZA 存储访问

可以使用 MSR/MRS 指令操作流式向量控制寄存器（SVCR） 来设置和读取 PSTATE.{SM, ZA} 位，具体操作如下：


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

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

• SMSTART：同时启用流式 SVE 模式和 ZA 存储访问
• SMSTART SM：启用流式 SVE 模式
• SMSTART ZA：启用 ZA 存储访问

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

• SMSTOP：同时禁用流式 SVE 模式和 ZA 存储访问
• SMSTOP SM：禁用流式 SVE 模式
• SMSTOP ZA：禁用 ZA 存储访问

下图显示了应用程序如何在流式 SVE 模式和非流式 SVE 模式之间切换：

有关使用 SMSTART 和 SMSTOP 在流式 SVE 模式和非流式 SVE 模式之间切换的更多信息，请参阅 Arm A-profile 架构参考手册的 C6.2.327 和 C6.2.328 节。

4. SME 架构状态

与 SVE2 类似，在流式 SVE 模式下，它具有 Z0-Z31 向量寄存器和 P0-P15 谓词寄存器。

编号最小的 SVE 向量寄存器 Zn 也包含固定长度的 Vn、Qn、Dn、Sn、Hn 和 Bn 寄存器。

当进入流式 SVE 模式（PSTATE.SM 从 0 变为 1）或退出流式 SVE 模式（PSTATE.SM 从 1 变为 0）时，所有这些寄存器都将被清零。

大多数非流式 SVE2 指令可以在流式 SVE 模式下使用，但可能使用不同的向量长度（流式模式使用 VSL 长度，非流式模式使用 NVSL 长度）。可以使用 RDSVL 指令读取当前有效的向量长度 VL。


   // 读取流式 SVE 向量寄存器大小的倍数到 Xd
RDSVL <Xd>, #<imm>

注意

由于 SME 支持向量长度不可知（VLA），在流式 SVE 模式下，软件很少需要显式读取 SVL 向量长度。在非流式 SVE 模式下，通常使用 RDSVL 指令来确定 SVL 的值。

5. ZA 存储阵列

SME 中新引入的 ZA（Z 阵列，ZA 存储） 是一个二维（2D）正方形阵列，大小为 SVL x SVL。它之所以被称为 Z 阵列，是因为其行和列的长度与流式 SVE 模式下的 Zn 寄存器长度一致。

例如：如果流式 SVE 模式下的向量长度为 256 位，即 Zn 寄存器的长度为 256 位，那么 ZA 的大小为 (256/8) 字节 x (256/8) 字节。

ZA 存储阵列可以通过以下方式访问：

• ZA 阵列向量访问
• ZA 分片
• ZA 分片子片

5.1 ZA 阵列向量访问

ZA 阵列的一行可以作为一个长度为 SVL 的向量来访问，该向量可以包含数据类型的元素长度为 8 位、16 位、32 位、64 位或 128 位，例如 32 位的 fp32 浮点数。


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

其中，B、H、S、D、Q 分别代表 8 位、16 位、32 位、64 位、128 位。

ZA 阵列向量的数量与 SVL 的字节数相同。例如，如果 SLV 为 256 位，则 ZA 阵列向量的数量为 32，N 的范围是从 0 到 31。

为了支持上下文切换，SME 引入了新的 LDR 和 STR 指令，用于从内存加载和存储一个 ZA 阵列向量。


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

5.2 ZA 分片

ZA 分片是 ZA 内的一个正方形二维子矩阵。一个 ZA 分片的宽度始终是 SVL，与 ZA 阵列的宽度相同。

ZA 可以被划分为多少个可用的 ZA 分片取决于元素的数据类型大小：

元素数据类型大小	分片数量	分片名称
8-bit	1	`ZA0.B`
16-bit	2	`ZA0.H` -`ZA1.H`
32-bit	4	`ZA0.S` -`ZA3.S`
64-bit	8	`ZA0.D` -`ZA7.D`
128-bit	16	`ZA0.Q` -`ZA15.Q`

• 当元素数据类型为 8 位时，ZA 只能作为一个 ZA 分片（ZA0.B）访问。
• 当元素数据类型为 16 位时，ZA 可以作为 2 个 ZA 分片（ZA0.H 和 ZA1.H）访问。
• 当元素数据类型为 32 位时，ZA 可以作为 4 个 ZA 分片（ZA0.S 到 ZA3.S）访问。
• 当元素数据类型为 64 位时，ZA 可以作为 8 个 ZA 分片（ZA0.D 到 ZA7.D）访问。
• 当元素数据类型为 128 位时，ZA 可以作为 16 个 ZA 分片（ZA0.Q 到 ZA15.Q）访问。

例如，如果 SVL 为 256 位，元素数据类型大小为 8 位，则 ZA 可以视为 ZA0.B，或者可以视为 32 个向量（32 行，每行大小为 32 x 8 位，即每行 32 个元素）。

如果 SVL 为 256 位，元素数据类型大小为 16 位，则 ZA 可以视为 2 个 ZA 分片（ZA0.H 和 ZA1.H），每个分片可以视为 16 个向量（16 行，每行大小为 16 x 16 位，即每行 16 个元素）。

这样做的优势是能够充分利用 ZA 存储。在实际应用中，例如，当 SVL 为 256 位，元素数据类型大小为 32 位，ZA 的大小为 256 位 x 256 位时，对两个 Z 寄存器中的向量执行外积运算，外积结果是一个 8 x 8 浮点数的 2D 数组。此外积仅需要 ZA 存储空间的 1/4。通过将 ZA 划分为 4 个 ZA 分片，可以充分利用 ZA 存储。

5.3 ZA 分片子片

一个 ZA 分片可以作为一个整体访问，也可以以单个 ZA 分片子片的形式访问。

   
    • 当作为整体访问时，指令可以使用分片的名称进行访问：
    
      ZA0.B, ZA0.H-ZA1.H, ZA0.S-ZA3.S, ZA0.D-ZA7.D 或 ZA0.Q-ZA15.Q

• ZA 分片子片是由其 ZA 分片在水平或垂直方向上的连续元素组成的一维数组，即 ZA 分片中的一行或一列。

访问 ZA 分片的一个向量就是读写一个 ZA 分片子片：

• 水平或垂直的 ZA 分片子片访问由 ZA 分片名称后的 H 或 V 后缀表示。
• 特定的 ZA 分片子片由一个索引表示，由 ZA 分片名称后的子片索引 [N] 指示。

例如，如果 SVL 为 128 位，元素数据类型大小为 8 位，则其水平和垂直 ZA 分片子片可以表示如下：

例如，如果 SVL 为 128 位，元素数据类型大小为 16 位，则其水平和垂直 ZA 分片子片可以表示如下：

为了提高硬件访问 ZA 分片和 ZA 分片子片的效率，一个 ZA 分片的 ZA 分片子片是交错排列的。

下图显示了这种交错排列的一个例子。在这个例子中，SVL 为 256 位，元素数据类型大小为 16 位。这意味着 ZA 可以被视为两个 ZA 分片（ZA0H 和 ZA1H），并且具有交错的水平分片子片：

下图显示了针对不同元素数据类型的水平和垂直 ZA 分片子片大小的混合视图：

左列显示了 ZA 内存每一行的不同处理方式。

• 设 SIZE 为向量元素的大小，其中 SIZE 为 1、2、4、8、16，分别代表数据类型 B、H、S、D 或 Q。
• 设 NUM_OF_ELEMENTS 为向量中元素的数量，即 bytes_of(SVL)/SIZE。
• 水平分片子片，ZAnH.<B|H|S|D|Q>[m] 访问一个向量，该向量包含 ZA 存储中的整行（m x SIZE + n）。该向量包含数据类型为 B、H、S、D 或 Q 的元素。
• 垂直分片子片，ZAnV.<B|H|S|D|Q>[m] 访问一个向量，该向量包含 ZA 存储中的整列（m x SIZE）。该向量包含数据类型为 B、H、S、D 或 Q 的元素。

ZAnV.[m] 访问一个包含列（m x SIZE）和行元素（i x SIZE + n）的向量，其中 i 的范围从 0 到 NUM_OF_ELEMENTS-1。该向量包含数据类型为 B、H、S、D 或 Q 的元素。

在应用混合元素数据类型大小以及水平和垂直分片子片时，请注意重叠问题。

有关 ZA 存储阵列、ZA 阵列向量、分片和分片子片的更多信息，请参阅 Arm A-profile 架构参考手册的 B1.4.8 至 B1.4.12 节。

6. 流式 SVE 模式下支持的指令

某些指令在流式 SVE 模式下受到限制：

• 一些 SVE/SVE2 指令变为非法执行
• 聚集加载和分散存储指令
• 使用 SVE2 的 首故障寄存器 指令
• 大多数 NEON 指令变为 UNDEFINED

有关受流式 SVE 模式影响的指令的更多信息，请参阅《Arm 架构参考手册》文档。

SME 增加了几条新指令，包括：

• 矩阵外积以及累加或相减指令，包括 FMOPA、UMOPA 和 BFMOPA。
• SVE2 向量寄存器（Z0-Z31）作为外积运算的行和列输入。
• ZA 存储存储二维矩阵分片的输出结果。
• 执行 SVE2 Z 向量与 ZA 的行或列进行加法运算的指令。
• 清除 ZA 分片的指令。
• 增加了一些可在流式和非流式模式下使用的指令。

7. SME 指令

操作 ZA 存储的主要 SME 指令包括：

• 计算两个向量的外积，然后累加或相减，并将结果放入 ZA 分片的指令。
• 将 SVE 向量（Z 寄存器） 存储或加载到 ZA 分片的行或列的指令。
• 在水平或垂直方向上，SVE 向量与 ZA 分片的加法指令。
• 将流式 SVE 模式的向量长度的倍数加到标量寄存器的指令。

7.1 外积及累加或相减指令

为了帮助理解外积及累加或相减指令，我们来看看如何利用外积运算来执行矩阵乘法。

计算两个向量 a 和 b 的外积将得到一个包含此外积的结果矩阵 C：

现在考虑两个矩阵 a 和 b 的矩阵乘法运算：

这个矩阵乘法可以通过计算两个外积运算并将两个结果矩阵累加来实现（这是常用手算方法），如下图所示：

SME 针对以下数据类型引入了高效的外积及累加或相减指令：

• 8 位、16 位整数
• FP16、BF16、FP32 和 FP64 浮点数

这些指令计算两个 Z 向量寄存器（Zn 和 Zm）中两个向量的外积，将得到的数组与 ZA 分片（ZAda） 中的现有数据累加或相减，并将结果存储在同一 ZA 分片（ZAda） 中。每个源向量由相应的控制谓词寄存器（Pn 和 Pm）独立地进行谓词控制。

输出数组	输入向量	描述	示例
`INT32`	`INT8` , `INT8`	将四个 `INT8` 的外积之和存入每个 `INT32` 元素	`SMOPA` 或 `SMOPS` 或 `UMOPA` 或 `UMOPS`：有符号或无符号整数外积和，及累加或相减。例如：`UMOPS <ZAda>.S, <Pn>/M, <Pm>/M, <Zn>.B, <Zm>.B`
`INT32`	`INT16` , `INT16`	将两个 `INT16` 的外积之和存入每个 `INT32` 元素	`SMOPA` 或 `SMOPS` 或 `UMOPA` 或 `UMOPS`：有符号或无符号整数外积和，及累加或相减。例如：`UMOPS <ZAda>.S, <Pn>/M, <Pm>/M, <Zn>.H, <Zm>.H`
`INT64`	`INT16` , `INT16`	如果实现了 `FEAT_SME_I16I64`，将四个 `INT16` 的外积之和存入每个 `INT64` 元素	`SMOPA` 或 `SMOPS` 或 `UMOPA` 或 `UMOPS`：有符号或无符号整数外积和，及累加或相减。例如：`UMOPS <ZAda>.D, <Pn>/M, <Pm>/M, <Zn>.H, <Zm>.H`
`FP32`	`BF16` , `BF16`	将两个 `BF16` 的外积之和存入每个 `FP32` 元素	`BFMOPA` 或 `BFMOPS`：`BFloat16` 外积和，及累加或相减。例如：`BFMOPS <ZAda>.S, <Pn>/M, <Pm>/M, <Zn>.H, <Zm>.H`
`FP32`	`FP16` , `FP16`	将两个 `FP16` 的外积之和存入每个 `FP32` 元素	`FMOPA` 或 `FMOPS`：半精度浮点数外积和，及累加或相减。例如：`FMOPS <ZAda>.S, <Pn>/M, <Pm>/M, <Zn>.H, <Zm>.H`
`FP32`	`FP32` , `FP32`	简单的 `FP32` 外积	`FMOPA` 或 `FMOPS`：浮点数外积，及累加或相减。例如：`FMOPS <ZAda>.S, <Pn>/M, <Pm>/M, <Zn>.S, <Zm>.S`
`FP64`	`FP64` , `FP64`	如果实现了 `FEAT_SME_F64F64`，执行简单的 `FP64` 外积	`FMOPA` 或 `FMOPS`：浮点数外积，及累加或相减。例如：`FMOPS <ZAda>.D, <Pn>/M, <Pm>/M, <Zn>.D, <Zm>.D`

7.1.1 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 的外积，外积的结果是图中的灰色矩阵，然后将此外积结果与 ZA 分片ZAda.S 中的现有值累加或相减，并将结果存储在同一 ZA 分片中。

7.1.2 FP16, BF16, INT16, INT8, I16I64 类型外积及累加或相减指令

由于这些指令会扩展计算结果的数据类型，这些操作不像前面的 FP32 和 FP64 类型指令那样直接。

• BF16 指令计算两个 BF16 的外积，将结果类型扩展为 FP32，然后与目标分片进行破坏性加法或减法。
• INT8 指令计算四个 INT8 的外积之和，将结果类型扩展为 INT32，然后与目标分片进行破坏性加法或减法。
• INT16 指令计算两个 INT16 的外积之和，将结果类型扩展为 INT32，然后与目标分片进行破坏性加法或减法。
• FP16 指令计算两个 FP16 的外积之和，将结果类型扩展为 FP32，然后与目标分片进行破坏性加法或减法。
• 如果实现了 FEAT_SME_I16I64，I16I64 指令计算四个 INT16 的外积之和，将结果类型扩展为 INT64，然后与目标分片进行破坏性加法或减法。

以下示例演示了 SVL 向量长度为 128 的 INT8UMOPA 指令的操作：


   UMOPA <ZAda>.S, <Pn>/M, <Pm>/M, <Zn>.B, <Zm>.B

每个输入寄存器（Zn.B, Zm.B）被视为包含 4x4 元素的矩阵，可以看作是已转置的、由 4 个连续元素组成的块（如图中红线标记所示）。

在这个例子中，由于 SVL 向量长度为 128 位：

• 第一个源向量 Zn.B 包含一个无符号 8 位整数的 4x4 子矩阵。
• 第二个源向量 Zm.B 包含一个无符号 8 位整数的 4x4 子矩阵。
• UMOPA 指令计算 4x4 扩展的 32 位整数外积之和，然后与目标分片（ZAda）中的整数进行破坏性累加。

更一般地说，UMOPA 指令将第一个源向量中的子矩阵与第二个源向量中的子矩阵相乘。每个源向量包含一个大小为 (SVL/32) x 4 的无符号 8 位整数的子矩阵。然后将得到的 (SVL/32) x (SVL/32) 扩展的 32 位整数外积破坏性地添加到一个 32 位整数目标分片。

以下示例演示了 SVL 为 128 位的 BF16BFMOPA 的操作：


   BFMOPA <ZAda>.S, <Pn>/M, <Pm>/M, <Zn>.H, <Zm>.H

在这个例子中，由于 SVL 向量长度为 128 位：

• 第一个源向量 Zn.H 包含一个 BF16 整数的 4x2 子矩阵，该子矩阵被扩展为单精度浮点数。
• 第二个源向量 Zm.H 包含一个 BF16 整数的 2x4 子矩阵，该子矩阵被扩展为单精度浮点数。
• BMOPA 指令计算一个 4x4 单精度外积之和，然后与目标分片（ZAda）中的单精度浮点数进行破坏性累加。

更一般地说，BFMOPA 指令将存储在第一个源向量中的 (SVL/32) x2BF16 子矩阵的类型扩展为单精度，将存储在第二个源向量中的 2x (SVL/32)BF16 子矩阵的类型扩展为单精度，并将这两个子矩阵相乘。然后，将得到的 (SVL/32) x (SVL/32) 单精度外积破坏性地添加到一个单精度目标分片。

下表显示了针对几种数据类型和 SVL 长度，执行一次外积及累加或相减指令所对应的 MACs（乘积累加） 数量：

	128-bit	256-bit	512-bit
`FP32`	16	64	256
`FP64`	4	16	64
`INT8`	64	256	1024
`INT16`	32	128	512
`BF16`	32	128	512
`FP16`	32	128	512

7.2 带谓词的 SME 指令

每个源向量都可以由其相应的控制谓词寄存器独立地进行谓词控制：

• 外积及累加或相减指令使用 Pn/M 和 Pn/M（没有 /Z 形式）：非活动的源元素被视为值为 0。
• 分片子片移动指令使用 Pg/M：目标子片中的非活动元素保持不变。
• 分片子片加载指令使用 Pg/Z：目标分片子片中的非活动元素设置为 0。
• 分片子片存储指令使用 Pg：非活动元素不会被写入内存。

谓词使得处理矩阵维度不是 SVL 倍数的情况更加容易。

例如，下图中指令：

输入向量 Z0 由 P0 谓词控制，Z1 由 P1 谓词控制。

在这个例子中：

• SVL 向量长度为 512 位。
• Z 寄存器包含一个由 16 个 FP32 数组成的向量。
• P0 中的最后两个元素是非活动的。
• P1 中的最后一个元素是非活动的。
• 该指令更新 ZA0.S 中的 (16-2) x (16-1) 个 FP32 元素，因为使用了 Pn/M，ZA0.S 中的其余元素保持不变。

下图显示了更多带谓词的外积及累加或相减的例子。图中带下划线的文本表示受非活动谓词元素影响的计算部分。

7.3 ZA 分片与 Z 向量的加法运算

SME 包含将向量与 ZA 分片的行或列相加的指令，这些指令也支持谓词。

指令	描述
`ADDHA`	将源向量加到 ZA 分片的每个水平子片上
`ADDVA`	将源向量加到 ZA 分片的每个垂直子片上

例如：


   ADDHA ZA0.S, P0/M, P1/M, Z1.S

将执行以下操作：

这条 ADDHA 指令将源向量 Z1 的每个元素加到 ZA0.S 分片的每个水平子片的相应活动元素上。

分片中的元素由一对主导谓词进行谓词控制。水平子片中的一个元素在以下条件下可被视为活动的：

• 它在第二个主导谓词的对应元素处为 TRUE，并且
• 它在第一个主导谓词的水平子片的行号处对应为 TRUE，并且目标分片中的非活动元素保持不变。

7.4 分片加载、存储、移动指令

SME 分片加载、存储、移动指令可以：

• 从内存读取数据并放入 ZA 分片的行或列
• 将 ZA 分片的行或列写入内存
• 将 ZA 分片的行移动到 SVE Z 向量寄存器
• 将 SVE Z 向量寄存器移动到 ZA 分片的行或列

7.4.1 分片子片加载和存储指令

LD1B、LD1H、LD1S、LD1D 和 LD1Q 指令分别将连续的内存值加载到包含 8 位、16 位、32 位、64 位或 128 位元素的 ZA 分片子片中。

ST1B、ST1H、ST1S、ST1D 和 ST1Q 指令分别将包含 8 位、16 位、32 位、64 位或 128 位元素的 ZA 分片子片存储到连续的内存中。

这些指令也支持谓词，例如：


   LD1B ZA0H.B[W0, #imm], P0/Z, [X1, X2]

这条 LD1B 指令执行带谓词的连续字节读取，将内存地址 (X1+X2) 处的数据读入 ZA0 中行号为 (W0+imm) 的水平分片子片。目标分片子片中的非活动元素设置为 0。


   ST1H ZA1V.H[W0, #imm], P2, [X1, X2, LSL #1]

这条 ST1H 指令执行带谓词的连续半字存储操作，将 ZA1 中列号为 (W0+imm) 的垂直分片子片存储到内存地址 (X1+X2*2) 处，分片子片中非活动的元素不会被写入内存。

7.4.2 分片子片移动指令

MOV 指令（MOVA 指令的别名）将 Z 向量寄存器的值移动到 ZA 分片子片，或者将 ZA 分片子片的值移动到 Z 向量寄存器。该指令对具有指定元素大小的 ZA 分片的单个水平或垂直子片进行操作。子片的行号/列号由子片的提取寄存器加上一个立即数偏移量指定。目标子片中的非活动元素保持不变。

例如：


   MOV     ZA0H.B[W0, #imm],  P0/M, Z0.B

或者


   MOVA  ZA0H.B[W0, #imm],  P0/M, Z0.B

该指令将向量寄存器 Z0.B 中的值移动到水平 ZA 分片子片ZA0H.B[W0,#imm]，使用 P0 作为谓词寄存器。目标分片子片中的非活动元素保持不变。

7.5 ZA 阵列向量加载/存储指令

SME LDR 指令从内存读取数据到 ZA 阵列向量，SME STR 指令将 ZA 阵列向量中的值存储到内存。这些指令不具备谓词功能。它们主要用于软件上下文切换期间保存/恢复 ZA 存储。当 PSTATE.ZA 启用时，SME LDR/STR 指令也可以在非流式 SVE 模式下使用。例如，以下 STR 指令中的 ZA 阵列向量由向量选择寄存器 Wv（标量寄存器 W）加上一个可选的立即数（Wv+Imm）指定。访问内存的地址是：一个标量寄存器作为基址，加上相同的可选立即数偏移量乘以当前向量长度的字节数。


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

7.6 ZA 分片清零指令

SME ZERO 指令可以清除一组 64 位的 ZA 分片：


   ZERO { <mask>}

ZERO 指令可以将最多 8 个名为 ZA0.D 到 ZA8.D 的 ZA 分片清零。要清零的分片由指令中的掩码指定，而其余分片保持不变。

当 PSTATE.ZA 启用时，该指令也可以在非流式 SVE 模式下使用。

如果要清除整个 ZA 阵列，可以使用指令别名 ZERO {ZA}。

7.7 新的 SVE2 指令

SME 架构扩展增加了一些新的 SVE2 指令，这些指令在非流式 SVE 模式下，也可以在实现了 SVE2 的 PE 中使用。这些指令包括：

• 选择谓词寄存器或全假的谓词选择指令
• 反转 64 位双字元素指令
• 有符号/无符号钳位到较小/较大值的向量指令

下面介绍谓词选择指令。

7.7.1 PSEL 指令

PSEL 指令选择一个谓词寄存器或全假到目标谓词寄存器，如下所示：


   PSEL <Pd>, <Pn>, <Pm>.<T>[<Wv>, <imm>]

如果第二个源谓词寄存器（Pm）中指定的元素为 True，则该指令将第一个源谓词寄存器（Pn）的内容放入目标谓词寄存器（Pd）中；否则，它将目标谓词寄存器的值设置为全假。例如，以下指令，假设 W12 的值为 0：


   PSEL P0, P1, P2.B[W12, #0]

第二个源谓词寄存器 [W12+0] 的第 [0] 个元素是 False，因此目标寄存器 P0 被设置为全 0（全假），如下图所示：

现在看下面的指令，仍然假设 W12 的值为 0，但这次立即偏移量为 1：


   PSEL P0, P1, P2.B[W12, #1]

第二个源谓词寄存器 [W12+1] 的第 [1] 个元素是 True，因此选择第一个源谓词寄存器的值到目标寄存器 P0，如下图所示：

8. SME C Intrinsic 编程

Arm 为 SME 提供了 C Intrinsic 函数，使开发者能够在 C/C++ 代码中直接使用 SME 指令，而无需编写汇编代码。这些 intrinsic 函数在 arm_sme.h 头文件中定义。

8.1 主要 Intrinsic 函数类别

8.1.1 流式 SVE 模式控制


   // 启用流式 SVE 模式和 ZA 存储
void svwrite_ver_za128(void);
void svwrite_ver_za256(void); 
void svwrite_ver_za512(void);
void svwrite_ver_za1024(void);
void svwrite_ver_za2048(void);

// 禁用流式 SVE 模式和 ZA 存储
void svzero_za(void);

8.1.2 外积运算


   // FP32 外积及累加
svfloat32_t svmopa_f32_m(
    svbool_t pg, svfloat32_t za, svfloat32_t zn, svfloat32_t zm);

// BF16 外积及累加  
svfloat32_t svmopa_bf16_m(
    svbool_t pg, svfloat32_t za, svbfloat16_t zn, svbfloat16_t zm);

// FP16 外积及累加
svfloat32_t svmopa_f16_m(
    svbool_t pg, svfloat32_t za, svfloat16_t zn, svfloat16_t zm);

// INT8 外积及累加
svint32_t svmopa_s8_m(
    svbool_t pg, svint32_t za, svint8_t zn, svint8_t zm);

8.1.3 ZA 存储访问


   // 加载 ZA 阵列向量
void svld1_ver_za8(uint32_t slice_base, svbool_t pg, void *ptr);
void svld1_ver_za16(uint32_t slice_base, svbool_t pg, void *ptr);
void svld1_ver_za32(uint32_t slice_base, svbool_t pg, void *ptr);
void svld1_ver_za64(uint32_t slice_base, svbool_t pg, void *ptr);

// 存储 ZA 阵列向量
void svst1_ver_za8(uint32_t slice_base, svbool_t pg, void *ptr);
void svst1_ver_za16(uint32_t slice_base, svbool_t pg, void *ptr);
void svst1_ver_za32(uint32_t slice_base, svbool_t pg, void *ptr);
void svst1_ver_za64(uint32_t slice_base, svbool_t pg, void *ptr);

8.1.4 分片子片操作


   // 移动数据到 ZA 分片子片
void svwrite_hor_za8_s8(uint32_t slice_base, svbool_t pg, svint8_t data);
void svwrite_hor_za16_s16(uint32_t slice_base, svbool_t pg, svint16_t data);
void svwrite_hor_za32_s32(uint32_t slice_base, svbool_t pg, svint32_t data);

// 从 ZA 分片子片读取数据
svint8_t svread_hor_za8_s8(uint32_t slice_base, svbool_t pg);
svint16_t svread_hor_za16_s16(uint32_t slice_base, svbool_t pg);
svint32_t svread_hor_za32_s32(uint32_t slice_base, svbool_t pg);

9. 完整矩阵乘法示例

以下是一个使用 SME C Intrinsic 实现单精度浮点数矩阵乘法的完整示例：


   #include <arm_sme.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

// 矩阵维度
#define M 64
#define N 64  
#define K 64

// 对齐内存分配
void* aligned_alloc(size_t alignment, size_t size) {
    void *ptr;
    if (posix_memalign(&ptr, alignment, size) != 0) {
        return NULL;
    }
    return ptr;
}

// 初始化矩阵
void init_matrix(float *matrix, int rows, int cols) {
    for (int i = 0; i < rows * cols; i++) {
        matrix[i] = (float)rand() / RAND_MAX;
    }
}

// 打印矩阵
void print_matrix(float *matrix, int rows, int cols) {
    for (int i = 0; i < rows; i++) {
        for (int j = 0; j < cols; j++) {
            printf("%8.4f ", matrix[i * cols + j]);
        }
        printf("\n");
    }
}

// 标准矩阵乘法 (参考实现)
void matrix_multiply_naive(float *A, float *B, float *C, int m, int n, int k) {
    for (int i = 0; i < m; i++) {
        for (int j = 0; j < n; j++) {
            float sum = 0.0f;
            for (int p = 0; p < k; p++) {
                sum += A[i * k + p] * B[p * n + j];
            }
            C[i * n + j] = sum;
        }
    }
}

// SME 加速矩阵乘法
void matrix_multiply_sme(float *A, float *B, float *C, int m, int n, int k) {
    // 启用 SME 流式模式 (假设 512-bit SVL)
    svwrite_ver_za512();
    
    // 获取 SVL 信息
    uint64_t svl = svcntb() * 4; // SVL 字节数
    
    // 计算每个维度需要处理的块数
    int block_m = (m + svl - 1) / svl;
    int block_n = (n + svl - 1) / svl;
    int block_k = (k + svl - 1) / svl;
    
    // 创建全真谓词
    svbool_t pg = svptrue_b32();
    
    for (int bm = 0; bm < block_m; bm++) {
        for (int bn = 0; bn < block_n; bn++) {
            // 清零当前输出分片
            svfloat32_t za_tile = svdup_f32(0.0f);
            
            for (int bk = 0; bk < block_k; bk++) {
                // 计算当前块的实际边界
                int m_start = bm * svl;
                int m_end = (bm + 1) * svl > m ? m : (bm + 1) * svl;
                int m_size = m_end - m_start;
                
                int n_start = bn * svl;
                int n_end = (bn + 1) * svl > n ? n : (bn + 1) * svl;
                int n_size = n_end - n_start;
                
                int k_start = bk * svl;
                int k_end = (bk + 1) * svl > k ? k : (bk + 1) * svl;
                int k_size = k_end - k_start;
                
                // 为当前块创建谓词
                svbool_t pg_m = svwhilelt_b32(m_start, m_end);
                svbool_t pg_n = svwhilelt_b32(n_start, n_end);
                svbool_t pg_k = svwhilelt_b32(k_start, k_end);
                
                // 加载 A 的当前块 (行向量)
                svfloat32_t a_block = svld1_vnum(pg_m, &A[m_start * k + k_start], 0);
                
                // 加载 B 的当前块 (列向量)  
                svfloat32_t b_block = svld1_vnum(pg_n, &B[k_start * n + n_start], 0);
                
                // 执行外积累加运算
                za_tile = svmopa_f32_m(pg, za_tile, a_block, b_block);
            }
            
            // 将结果存储回内存
            int c_start_row = bm * svl;
            int c_start_col = bn * svl;
            int c_rows = ((bm + 1) * svl > m) ? (m - bm * svl) : svl;
            int c_cols = ((bn + 1) * svl > n) ? (n - bn * svl) : svl;
            
            // 使用水平分片子片存储结果
            for (int i = 0; i < c_rows; i++) {
                svbool_t row_pg = svwhilelt_b32(0, c_cols);
                svfloat32_t row_data = svread_hor_za32_s32(i, row_pg);
                svst1_vnum(row_pg, &C[(c_start_row + i) * n + c_start_col], 0, row_data);
            }
        }
    }
    
    // 禁用 SME 模式
    svzero_za();
}

// 验证结果
int verify_results(float *C1, float *C2, int m, int n, float tolerance) {
    for (int i = 0; i < m * n; i++) {
        if (fabs(C1[i] - C2[i]) > tolerance) {
            printf("Mismatch at element %d: %f vs %f\n", i, C1[i], C2[i]);
            return 0;
        }
    }
    return 1;
}

int main() {
    printf("SME Matrix Multiplication Example\n");
    
    // 分配对齐内存
    float *A = aligned_alloc(64, M * K * sizeof(float));
    float *B = aligned_alloc(64, K * N * sizeof(float));
    float *C_naive = aligned_alloc(64, M * N * sizeof(float));
    float *C_sme = aligned_alloc(64, M * N * sizeof(float));
    
    if (!A || !B || !C_naive || !C_sme) {
        printf("Memory allocation failed\n");
        return 1;
    }
    
    // 初始化输入矩阵
    srand(42);
    init_matrix(A, M, K);
    init_matrix(B, K, N);
    memset(C_naive, 0, M * N * sizeof(float));
    memset(C_sme, 0, M * N * sizeof(float));
    
    printf("Matrix dimensions: %d x %d x %d\n", M, K, N);
    
    // 执行标准矩阵乘法
    printf("Running naive matrix multiplication...\n");
    matrix_multiply_naive(A, B, C_naive, M, N, K);
    
    // 执行 SME 加速矩阵乘法
    printf("Running SME accelerated matrix multiplication...\n");
    matrix_multiply_sme(A, B, C_sme, M, N, K);
    
    // 验证结果
    printf("Verifying results...\n");
    if (verify_results(C_naive, C_sme, M, N, 1e-5f)) {
        printf("✓ Results match!\n");
    } else {
        printf("✗ Results don't match!\n");
    }
    
    // 打印部分结果用于验证
    printf("\nFirst 4x4 elements of result matrix:\n");
    print_matrix(C_sme, 4, 4);
    
    // 释放内存
    free(A);
    free(B);
    free(C_naive);
    free(C_sme);
    
    return 0;
}

编译说明

要编译上述代码，需要使用支持 SME 的 Arm 编译器：


   # 使用 Arm Compiler for Linux
armclang -march=armv9-a+sme -O3 -o sme_matmul sme_matmul.c

# 或者使用 GCC with SME support  
aarch64-none-linux-gnu-gcc -march=armv9-a+sme -O3 -o sme_matmul sme_matmul.c

示例说明

1. 流式模式管理：使用 svwrite_ver_za512() 启用 SME 流式模式，svzero_za() 禁用。
2. 分块处理：将大矩阵分解为适合 ZA 存储的小块进行处理。
3. 外积运算：使用 svmopa_f32_m() 执行高效的矩阵外积累加运算。
4. 谓词控制：使用 svwhilelt_b32() 创建适当的谓词来处理边界情况。
5. 数据移动：使用 svld1_vnum() 和 svst1_vnum() 在内存和寄存器间传输数据。

这个示例展示了如何利用 SME 的 ZA 存储和外积指令来高效实现矩阵乘法，特别适合处理大规模矩阵运算。

参考资料

Arm Scalable Matrix Extension (SME) Introduction^[1]
Arm Scalable Matrix Extension (SME) Instructions^[2]
Arm C Language Extensions for SME^[3]

引用链接

[1] Arm Scalable Matrix Extension (SME) Introduction: https://developer.arm.com/community/arm-community-blogs/b/architectures-and-processors-blog/posts/arm-scalable-matrix-extension-introduction
[2] Arm Scalable Matrix Extension (SME) Instructions: https://developer.arm.com/community/arm-community-blogs/b/architectures-and-processors-blog/posts/arm-scalable-matrix-extension-introduction-p2
[3] Arm C Language Extensions for SME: https://developer.arm.com/documentation/101726/0400

【声明】内容源于网络