Armv9 技术讲堂 | SME 详解

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

SME 显著提升了在 Arm CPU 上处理现有 AI 和 ML 应用的能力，从而让各类 AI 加持的设备和应用有更快、更及时的用户体验。它也可以加速很多利用矩阵运算的应用，例如 AR、VR 和图像处理，它们在 AI 和 ML 扮演着越来越重要的角色。

与 Arm CPU 可以处理各类神经网络中诸多数据类型类似，SME 提供了不断进化和日益复杂的 AI 和 ML 应用需要的灵活性。这保证了 Arm 架构在 AI 快速演化的年代，仍然承担 AI 计算重要职责。

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

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

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

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

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

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

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

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

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

在同一应用中很容易动态切换 Streaming SVE 和 Non-streaming SVE 模式。SME 引入了一些新的 PSTATE 位，PSTATE.{SM, ZA} 用于启用和禁用 Streaming SVE 模式和 SME ZA 存储：

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

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

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

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

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

物理上，SIMD (Neon)、浮点 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 长度的向量，与 NSVL 可能不一样。生效的 SVL 长度可以通过 RDSVL 指令读出来：

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

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

例如，如果 SVL 的长度为 256 位（32 字节），那么 Zn 寄存器的长度为 32 字节（256 位），ZA 的大小为 32 字节 x 32 字节（或 256 位 x 256 位）。

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

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

ZA array 向量的数量与 SVL 中的字节数相同，例如，如果 SVL 是 256 位（32 字节），那么 ZA array 向量的数量就是 32 个 (0<= N <32)。

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

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

这样做可以充分利用 ZA 存储，例如，如果元素数据类型为 32 位浮点数且 SVL 为 256 位（SVL 长度可以放 8 个 32 位浮点数），那么 ZA 的大小为 32 x 32 字节 （可以放 32 行 x 8 列的二维浮点数数值），一个 Z 寄存器可以放 8 个浮点数元素的向量。那么两个 Z 寄存器里的向量做外积运算，产生的外积结果是 8 行x 8 列的二维浮点数数组，这个外积只需要 1/4 的 ZA 存储。为了充分利用，这样 ZA 可以分成 4 个 ZA tiles，分别从 ZA0.S 到 ZA3.S。

再例如，如果 SVL 为 256 位（32 字节），元素的数据类型大小为 8 位，那么 ZA 可以看成为一个有 32 行，每行为 32 个 8 位的数组。

如果 SVL 同为 256 位（32 字节），但元素的数据类型大小为 16 位，那么 ZA 可以看出是两个 ZA tiles， ZA0.H 和 ZA1.H，每个 tile 由 16 行，每行有 16 个 16 位的数组组成。

一个 ZA tile 可以作为一个整体来访问，也可以以一个个 ZA tile slice 的方式访问。当作为一个整体访问时，指令可以使用 tile 的名字访问：

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

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

例如，如果 SVL 为 128 位 且元素类型的大小为 8 位，那么其水平和垂直的 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 位且元素类型的大小为 16 位，那么其水平和垂直的 ZA tile slices 可以由下图展示：

事实上，为了实现高效的硬件对 ZA tile 和 tile slice 的访问，ZA tile 的 tile slice 在 ZA 中以交织方式存储。例如，如果 SVL 是 256 位，元素数据类型大小是 16 位，那么其水平 ZA tile slice 在 ZA 中的存储可以展示为：

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

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

图中左边的那些列显示了 ZA 存储的每行可以被各种数据类型大小的 tile 水平 slice 访问的方式。 设 'SIZE' 为向量元素的类型大小，B、H、S、D 和 Q 类型的 SIZE 分别为 1、2、4、8、16。

设 'NUM_OF_ELEMENTS' 为可以放在向量里的元素个数，即（SVL 的字节数）/SIZE。那么，水平 tile slice，ZAnH.[m] 访问 ZA 存储的第 (m*SIZE+n) 行（完整一行）。例如，ZA2H.S[5] 访问 ZA 存储的第 (5*4+2) 行，即第 22 行。ZA6H.D[0] 访问 ZA 存储的第 (0*8+6) 行，即第 6 行。

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

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

有些 Neon/SVE2 指令在 Streaming SVE 模式受到影响：

SME 主要增加以下指令：

END