Armv9 技术讲堂 | SME 详解

Arm 技术专家将分上下两期带大家深入了解 SME 及其指令，为你详细解析 SME 如何确保 Arm CPU 能够高效地处理新的高能耗生成式 AI 工作负载，从而在数十亿台 Arm 技术驱动的设备上提供更好的 AI 体验！

本文引用地址：

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 的基础之上，新增了高效处理矩阵的能力。其关键功能包括：

计算两个 SVE 向量的外积 (outer product)

矩阵块 (tile) 的存储

存取矩阵块中的向量，向矩阵块中插入向量和提取矩阵块里的向量，包括 on-the-fly 矩阵转置

Streaming SVE 模式

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

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

新的架构状态 (architectural state)，可以用来存储二维矩阵块。

新的 Streaming SVE 模式支持运行 SVE2 指令，Streaming SVE 模式下，SVE2  向量具有矩阵块一样宽度，通常比 Non-streaming SVE 模式下的向量要更宽。

新的外积指令用于计算两个向量的外积，且可以进行外积的累加或累减，将结果放入一个矩阵块。

新的存取和移动 (load, store, move) 指令，用于从矩阵块的一行或是一列读取一个向量到向量寄存器，或是将一个向量写入矩阵块的一行或是一列。

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

Streaming SVE 模式

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

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

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

大多数新加入的 SME 指令只能在 Streaming SVE 模式下使用。Streaming SVE 模式下的硬件向量长度 (Streaming vector length, SVL) 可以和 Non-streaming SVE  模式的向量长度 (Non-streaming vector length, 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 模式

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

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

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

SM：启用和禁用 Streaming SVE 模式

ZA：启用和禁用 ZA 存储访问

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

MSR SVCRSM, #

MSR SVCRZA, #

MSR SVCRSMZA, #

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

SMSTART：启用 Streaming SVE 模式和 ZA 存储访问

SMSTAT SM：启用 Streaming SVE 模式

SMSTART ZA：启用 ZA 存储访问

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

SMSTOP：禁用 Streaming SVE 模式和 ZA 存储访问

SMSTOP SM：禁用 Streaming SVE 模式

SMSTOP ZA：禁用 ZA 存储访问

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

SME 架构状态

与 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 长度。

ZA array

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 可以通过以下方式访问：

ZA array 向量

ZA tiles

ZA tile slices

ZA array 向量访问

ZA array 的一行可以当成一个 SVL 长度的向量（ZA array 向量）来访问，这个向量可以放数据类型长度位 8 位、16 位、32 位、64 位、128 位的元素：

其中 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 向量：

LDR ZA[,], [{, #, MUL VL}]

STR ZA[, ], [{, #, MUL VL}]

ZA tiles

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

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

当元素数据类型为 8 位时，ZA 只能是 1 个 ZA tile，ZA0.B

当元素数据类型为 16 位时，ZA 可分为 2 个 ZA tiles，ZA0.H-ZA1.H

当元素数据类型为 32 位时，ZA 可分为 4 个 ZA tiles，ZA0.S-ZA3.S

当元素数据类型为 64 位时，ZA 可分为 8 个 ZA tiles，ZA0.D-ZA7.D

当元素数据类型为 128 位时，ZA 可分为 16 个 ZA tiles，ZA0.Q-ZA15.Q

这样做可以充分利用 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 可以作为一个整体来访问，也可以以一个个 ZA tile slice 的方式访问。当作为一个整体访问时，指令可以使用 tile 的名字访问：

一个 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 位 且元素类型的大小为 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 间的重叠。

Steaming SVE 模式下支持的指令

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

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

Gather/Scatter load/store SVE2 指令

使用 First Fault 寄存器的 SVE2 指令

大多的 Neon 指令变为 UNDEFINED

SME 主要增加以下指令：

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

SVE2 向量寄存器 (Z0-Z31) 作为外积运算的行和列输入

ZA 存储用来保存二维的矩阵块输出

将 SVE2 Z 向量与 ZA 的行或列做加法运算的指令

对 ZA tiles 的清零操作指令

增加了一些在 Streaming SVE 和 Non-streaming SVE 模式下都能使用的指令