Arm构架如何让AI应用高效运行于CPU (1)

引言

Arm为在各类技术领域运行无处不在的人工智能 (AI) 奠定了坚实的技术基础。Arm 始终专注于快速地进行架构演进，确保arm生态系统能够适应未来的技术趋势和不断变化的计算需求。过去二十多年来，Arm 一直致力于为 AI 创新奠定基础，先是在 Armv7 架构中引入了高级单指令多数据 (SIMD) 扩展，初探机器学习 (ML) 工作负载，再是如今的 Armv9 架构，arm在 Arm CPU 上集成了用于加速和保护如大语言模型 (LLM) 等先进的生成式 AI 工作负载的特性。

AI, ML的主要运算是乘加（MAC）和矩阵运算，arm A-profile的指令对其的支持和加速也在持续的进行。演化的方向主要包括：

1. 单条向量指令执行的乘加（MAC）的数量越来越多。其中包含：
   a. 从armv8.0-a NEON开始就支持的向量乘累加（MAL）操作，到armv8.4-a引入的向量点积（Dot product）运算，armv8.6-a引入的矩阵乘(MM)指令，和armv9-a SME引入的外积(Out Product)运算。这些操作还提供另一个好处：提高了一次读取的数据，在这些MAC运算中的利用次数，从而提供更高的（运算次数/内存访问）的比例。这对大矩阵处理的内存访问更加高效，减少内存访问和cache miss的次数。
   b. 更长的向量支持。SVE提供更长更加灵活的向量长度支持，SVE更长的向量提供了相对NEON成向量长度倍数增长的运算数量。Armv9-a SME更是提供了向量长度倍数平方（长度倍数*长度倍数）的运算数量的增长。SME2还提供了多向量（multiple vectors）运算指令，可以在一条指令中操作多个SVE向量寄存器进行运算。
   
2. 支持更加丰富的数据类型，依次加入了单精度浮点(FP16), BFloat16 (BF16)和8 bit 浮点(FP8)的支持，并可以利用MAL, Dot product, 矩阵乘(MM), SME out product来对它们进行运算。
   

通过arm构架文档来看懂这些指令，会有些困难。本文通过直观的图示和解释来帮助大家理解这些指令，从而使大家可以利用好这些指令来优化AI算法。

Armv8-a NEON 乘加指令

虽然目前很多平台的CPU已经使用了armv9-a构架的CPU，但还是有一些软件代码是基于armv8.0-a NEON优化的，通过理解这些指令可以帮助利用更好更新的向量指令来进一步优化这些代码。
NEON是固定的128-bit向量长度的，它提供了丰富的对整型，浮点型数据的乘累加，乘累减，饱和运算，乘并取半，各种round操作。这些指令对AI之前的DSP运算提供良好的支持。

Base浮点数乘加指令执行相对简单的向量乘累加操作，支持的数据类型包括单精度，双精度，半精度浮点数。支持的形式包括：一个向量与另一个向量的乘累加，和一个向量与另一个向量的一个元素进行乘累加操作。

• FMLA (vector): Floating-point fused multiply-add to accumulator (vector).
• FMLA (by element): Floating-point fused multiply-add to accumulator (by element).

以单精度浮点数据类型FMLA (vector)为例：

FMLA <Vd>.4S, <Vn>.4S, <Vm>.4S

它可以在128-bit长度的向量寄存器里存放4个FP32数据组成的向量，每条lane上执行如下Op操作（图示了每个Op进行的操作）：

每个Op操作执行一次FP32乘法＋一次累加操作。那么一条FMLA(vector)指令执行4xOp操作，也就是４次乘＋４次加（４FP32 MAC）。

FMLA 还支持(by element)的形式，其操作为：

FMLA <Vd>.4S, <Vn>.4S, <Vm>.S[Index]

指令中的Index用于指示使用的是第二个向量中的那一个元素参与运算。

FMLA (by element)指令也提供了４x FP32 MAC，这条指令也在矩阵乘法运算在比较常见。
为了保持运算的精度，NEON指令也提供了将乘加运算结果的类型变宽的操作。但是如果源操作数的向量长度是128-bit，而存储结果的寄存器长度也是固定的128-bit，那么要完成完整的128-bit源向量的运算则需要两条指令。支持这个操作的的有：

• FMLAL, FMLAL2：　Floating-point fused multiply-add long to accumulator
• FMLALB, FMLALT：　8-bit floating-point multiply-add long to half-precision
• FMLALLBB, FMLALLBT, FMLALLTB, FMLALLTT　：8-bit floating-point multiply-add long-long to single-precision
• BFMLALB, BFMLALT : BFloat16 floating-point widening multiply-add long

FMLAL(vector), FMLAL2(vector)的操作为：

FMLAL <Vd>.4S, <Vn>.4S, <Vm>.4H

这条指令将两个源NEON寄存器中低64-bit的对应半精度浮点值相乘，并将中间乘积（不进行舍入）累加到目标NEON寄存器中相应的单精度浮点向量元素中。

FMLAL2 <Vd>.4S, <Vn>.4S, <Vm>.4H

这条指令将两个源NEON寄存器中高64-bit的对应半精度浮点值相乘，并将中间乘积（不进行舍入）累加到目标NEON寄存器中相应的单精度浮点向量元素中。
其中Op的操作为：

通过组合FMLAL和FMAL2可以完成完整128-bit向量中半精度浮点数的乘加操作，并且结果类型扩大为单精度。

FMLAL和FMLAL2也支持by element的操作：

FMLAL2 <Vd>.4S, <Vn>.4S, <Vm>.H[Index]

上面的指令是以源向量的低64-bit和高 64-bit来分别运算的，另外一类指令是以源向量中元素的奇数或是偶数编号来分别运算的。以BFMLALB, BFMLALT为例（FMLALB, FMLALT，FMLALLBB, FMLALLBT, FMLALLTB, FMLALLTT类似）：

BFMLALB <Vd>.4S, <Vn>.8H, <Vm>.8H

这条指令将第一个和第二个源向量中奇数编号（Top）的16位元素从Bfloat16扩展为单精度格式。然后，该指令将这些值相乘并累加到目标向量中与源向量中相应BFloat16元素重叠的单精度元素中，而不进行中间舍入。

BFMLALT <Vd>.4S, <Vn>.8H, <Vm>.8H

这条指令将第一个和第二个源向量中偶数编号（Bottom）的16位元素从Bfloat16扩展为单精度格式。然后，该指令将这些值相乘并累加到目标向量中与源向量中相应BFloat16元素重叠的单精度元素中，而不进行中间舍入。
其中Op的操作为：

BFMLALB，BFMLALT也支持by element操做，即一个向量与两一个向量的一个元素乘加操作。
对于整型数据的支持，NEON支持64-bit, 32-bit, 16-bit和8-bit整型乘加运算，支持如上面的源和目的向量数据类型相同，或是扩大的操作。包括有符号数之间，无符号数之间的乘加操作。
简单的整型乘加（MLA）向量运算有MLA指令，可以执行两个向量之间的乘加操作，源和目的的数据类型大小相同。例如

MLA <Vd>.4S, <Vn>.4S, <Vm>.4S

其中8-bit的整型/定点数的支持，对AI推理运算特别有用。以有符号8-bit乘加结果类型扩大的UMLAL, UMLAL2为例：

UMLAL <Vd>.8H, <Vn>.16B, <Vm>.16B

这条指令将第一个源NEON寄存器的上半部分中的向量元素与第二个源NEON寄存器中对应的向量元素相乘，并将结果累加到目标NEON寄存器的向量元素中。目标向量元素的长度是相乘元素长度的两倍。

UMLAL2 <Vd>.8H, <Vn>.16B, <Vm>.16B

这条指令将第一个源NEON寄存器的下半部分中的向量元素与第二个源NEON寄存器中对应的向量元素相乘，并将结果累加到目标NEON寄存器的向量元素中。目标向量元素的长度是相乘元素长度的两倍.
其中的Op操作为：

这些指令可以利用128-bit的NEON向量寄存器，在保持乘加运算精度的同时，进行8次MAC操作。

Dot product（点积）指令

Dot product指令的引入进一步增加了单条指令中的乘加运算数量。Dot product刚开始时在armv8.4-a引入（ arm Armv8.2-a CPU如cortex-a55, a75等也支持这些指令 ），只支持int8类型的Dot product, 之后构架进一步丰富了Dot product支持的数据类型。SVE中的Dot Product的形式进一步丰富。

• SDOT (by element): Dot product signed arithmetic (vector, by element).
• SDOT (vector): Dot product signed arithmetic (vector).
• SUDOT (by element): Dot product with signed and unsigned integers (vector, by element).
• UDOT (by element): Dot product unsigned arithmetic (vector, by element).
• UDOT (vector): Dot product unsigned arithmetic (vector).
• USDOT (by element): Dot product with unsigned and signed integers (vector, by element).
• USDOT (vector): Dot product with unsigned and signed integers (vector).
• BFDOT (indexed): BFloat16 floating-point indexed dot product.
• BFDOT (vectors): BFloat16 floating-point dot product.
• FDOT (8-bit floating-point to half-precision, by element): 8-bit floating-point dot product to half-precision (vector, by element).
• FDOT (8-bit floating-point to half-precision, vector): 8-bit floating-point dot product to half-precision (vector).
• FDOT (8-bit floating-point to single-precision, by element): 8-bit floating-point dot product to single-precision (vector, by element).
• FDOT (8-bit floating-point to single-precision, vector): 8-bit floating-point dot product to single-precision (vector).

我们先看一下整型的Dot product，以Int8(包括有符号，无符号)类型的Dot product为例：

SDOT <Vd>.4S, <Vn>.16B, <Vm>.16B

这条指令对每个32位元素中的四个有符号8位元素进行点积运算，第一个源寄存器中的四个有符号8位元素与第二个源寄存器中对应32位元素的四个有符号8位元素相乘，并将结果累加到目标寄存器中对应的32位元素中。
它相当于对一个4x4的int8的矩阵A的行ｎ与另一个4x4的int8的矩阵B的列n进行点积运算，如图中的Op所示，Op做的操作为c0=c0+a0b0+a1b1+a2b2+a3b3。因此128-bit的vector寄存器中，这条Dot product指令进行4x(4次乘＋４次加)，即16x Int8 MAC，并且因为结果数据类型扩大，保持了精度。相对之前的 SMLAL, SMLAL2指令，单条指令的MAC数量增加一倍。

这类Dot Product也支持by element的形式：

SDOT <Vd>.4S, <Vn>.16B, <Vm>.4B[<index>]

这条指令对每个32位元素中的四个8位元素进行点积运算，将第一个源寄存器中的四个8位元素与第二个源寄存器中索引(Index)指定的32位元素的四个8位元素相乘，并将结果累加到目标寄存器中对应的32位元素中。

再以BF16的Dot product为例：

BFDOT <Vd>.4S, <Vn>.8H, <Vm>.8H

这条指令对源向量中每对相邻的BFloat16元素执行unfused的乘积求和操作。然后中间的求和结果在累加到目标向量中与源向量中对应的单精度元素。

BF16的Dot product也支持by element的形式：