Arm 始终专注于架构演进，确保生态系统能够适应未来的技术趋势和不断变化的计算需求。Armv9 架构上的可伸缩矩阵扩展 (SME) 显著提高了 Arm CPU 对现有人工智能 (AI) 和机器学习 (ML) 工作负载的处理能力，从而在各种 AI 驱动的设备和应用中带来速度更快、响应更灵敏的用户体验。
在此前的内容中，Arm 技术专家为大家简要介绍了 SME 和 SME 指令，本期将带你了解如何使用 Neon、可伸缩向量扩展 (SVE) 和 SME 这三种不同的 Arm 技术实现相同的矩阵-矩阵乘法算法。这三个示例展现了这些技术之间的关键差异，为开发者将代码从 Neon、SVE/SVE2 移植到 SME/SME2 提供指导。
架构演进
Armv7 引入了高级 SIMD 扩展，为一系列整型和浮点型提供单指令多数据 (SIMD) 操作。Neon 是高级 SIMD 指令的一种实现方案，作为部分 Arm Cortex-A 系列处理器的扩展提供。2011 年 Armv7-A 中引入了 Neon。Neon 提供固定宽度的 128 位寄存器。这意味着每条 Neon 指令对固定数量的数据值进行操作，例如四个 32 位数据值。
2016 年 Armv8-A 中引入了 SVE，2021 年 Armv9-A 中引入了 SVE2，它们提供可变长度寄存器。寄存器的大小由实现方案定义，从 128 位到 2048 位寄存器不等。这意味着程序员不知道可用寄存器的大小，因此必须将代码编写为与向量长度无关。因此，每条指令处理的数据值的数量不是固定的，而是可变的。
2021 年 Armv9-A 中引入了 SME 和 SME2，也提供可变长度寄存器。SME 引入了两个关键的新架构特性：Streaming SVE 模式和 ZA 存储。Streaming SVE 模式是一种高吞吐量的矩阵数据处理模式，ZA 存储则是一种专用的二维数组，可以方便地进行常见的矩阵操作。这些特性使 SME 和 SME2 能够高效地处理矩阵和基于向量的工作负载。
这些 SIMD 架构扩展提供的指令可加速多种应用，包括媒体和信号处理应用、高性能计算 (HPC) 应用，以及 ML 应用。
本文中的示例使用了内置函数 (intrinsics)，即编译器提供的与特定 Arm 指令相对应的函数。这使得程序员能够用 C 语言而不是汇编语言编写整个程序。
矩阵-矩阵乘法
本文中的所有三个示例都能够实现矩阵-矩阵乘法。矩阵乘法需要接收两个输入矩阵，通过将第一个矩阵一行的每个元素与第二个矩阵一列的相应元素相乘，然后对这些乘积求和，从而生成一个结果矩阵。结果矩阵的维度由第一个矩阵的行数和第二个矩阵的列数决定。例如，3 x 2 矩阵乘以 2 x 3 矩阵将得到 3 x 3 矩阵。
要将矩阵 A 和矩阵 B 相乘，矩阵 A 的列数必须等于矩阵 B 的行数。将矩阵 A 和矩阵 B 相乘将得到矩阵 C。
Neon
此示例使用 Neon 内置函数来执行矩阵-矩阵乘法。相关代码执行以下操作：
- 两个输入矩阵包含以列优先格式存储的 32 位浮点数据。
- 代码以 4 x 4 的块形式迭代处理这些矩阵中的所有数据。
- vld 内置函数将输入矩阵的行和列中的四个值加载到 Neon 寄存器中。
- 每个 fma Neon 内置函数执行四次乘加运算，计算正在处理的 4 x 4 块的结果。
- vst 内置函数将结果矩阵存储到内存中。
以下是使用 Neon 内置函数的示例代码：
此示例使用以下 Neon 代码特性：
此示例使用 4 x 4 的固定块大小。这意味着输入矩阵在两个维度上都必须是四的倍数。可以通过用零填充矩阵来处理其他大小的矩阵。
SVE/SVE2
此示例使用 SVE2 内置函数来执行矩阵-矩阵乘法。
Neon 示例和 SVE2 示例之间的主要区别在于 SVE2 使用可变长度向量。Neon 示例可以使用 4 x 4 的固定块大小来匹配 Neon 寄存器中的四个 32 位值，但程序员要到运行时才能知道 SVE2 寄存器的大小。这意味着代码必须与向量长度无关。此示例使用 predication 来控制 SVE2 内置函数操作的数据值的数量。这意味着无论实现的大小如何，它们都能够精准地适配 SVE2 寄存器。Neon 示例使用 32 位浮点数据类型 float32x4\_t，其中 的“4”表示每个 Neon 寄存器可以包含四个 32 位值。SVE2 示例使用 svfloat32\_t 数据类型，因为无法在运行前知道 SVE2 寄存器的大小。
相关代码执行以下操作：
- 两个输入矩阵包含以列优先格式存储的 32 位浮点数据。
- 代码以四行一组的形式迭代处理这些矩阵中的所有数据。它使用 svcntw 内置函数，返回向量中 32 位元素的数量，以匹配加载到 SVE2 寄存器大小的列数。这有助于避免对外循环每次迭代中的元素数量进行硬编码。whileit 内置函数生成一个 predicate，以确保不超过矩阵的界限。
- 四个 svld 内置函数使用之前生成的 predicate 将矩阵数据加载到 SVE2 寄存器中。
- svlma 内置函数执行乘加运算，计算当前迭代的结果。
- svst 内置函数将结果矩阵存储到内存中。
以下是使用 SVE2 内置函数的示例代码：
此示例使用以下 SVE2 代码特性：
SME/SME2
此示例使用 SME2 汇编指令来执行矩阵-矩阵乘法。SME2 示例与其他示例的区别如下：
- SME2 示例使用汇编代码，而不是其他示例所使用的内置函数。
- SME2 提供 ZA 存储，这是专为矩阵运算设计的二维数据数组。此 ZA 存储内的子数组可作为 tile 进行访问，并且 tile 内的元素可以垂直或水平访问。这为操作矩阵数据提供了非常灵活的机制。
- SME2 提供了执行矩阵运算的新指令。例如，fmopa 指令可计算外积。
- SME2 提供了一种多向量二维 predication 机制，以确保不超出矩阵边界。
Streaming SVE 模式（使用 smstart 指令进入的）启用 SME2 指令和 ZA 存储。
此示例使用 matLeft 和 matRight 作为输入矩阵。示例中用到了以下运算原理：两个矩阵相乘等同于依次对 matLeft 的每一列和 matRight 的每一行的外积求和。
初始输入矩阵作为行优先数组存储在内存中。矩阵乘法是将 matLeft 的一列与 matRight 的一行的外积求和。由于外积需要 matLeft 的列元素，因此代码重新排列 matLeft 数据，以便列元素连续存储在内存中。为了简洁起见，本文未展示此类数据重排，如需参考可查看 SME 程序员指南。
此示例包含三个嵌套循环：
- 最外层循环用于迭代处理结果矩阵的行。
- 中间层循环用于迭代处理结果矩阵的列。
- 最内层循环用于迭代处理 K 维度，通过对乘积求和生成结果矩阵元素。
使用 ld1w 指令将矩阵数据从内存加载到 ZA 存储。外积计算使用 fmopa 指令。每个 fmopa 指令读取两个 SVE Z 输入向量，并使用结果更新整个 SME ZA tile。二维 predication 确保不超出矩阵的边界。最后，st1w 指令将结果从 ZA 存储写入内存。
进一步阅读
有关 SME 的更多详细信息，请参照 Arm Architecture Reference Manual for A-profile Architecture 架构文档：https://developer.arm.com/doc...，也可点击阅读原文，即刻访问。
有关代码如何逐行运行的详细信息，请参阅 SME Programmer's Guide：https://developer.arm.com/doc...
本文为 Arm 原创文章，转载请留言联系获得授权并注明出处。
文章来源：Arm社区
推荐阅读
