修志龙_ZenonXiu · 9月6日 · 上海

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

矩阵乘(Matrix Multiply)指令

image.png
为了进一步利用向量寄存器,在向量运算中执行更多的乘加(MAC)操作。Armv8.6-a引入了矩阵乘(Matrix Multiply)指令。这些指令相当于取A矩阵的两行放在向量Vec_A中,取B矩阵的两列放在向量Vec_B中,矩阵乘指令执行:

  • 将Vec_A中第一行与Vec_B中的第一列做点积
  • 将Vec_A中第一行与Vec_B中的第二列做点积
  • 将Vec_A中第二行与Vec_B中的第一列做点积
  • 将Vec_A中第二行与Vec_B中的第二列做点积

从而得到2x2的结果矩阵。
NEON矩阵乘(Matrix Multiply)指令支持的数据类型包括整型,BF16浮点型的支持。SVE进一步加入了单精度,双精度的支持。NEON矩阵乘(Matrix Multiply)指令:

  • SMMLA (vector): Signed 8-bit integer matrix multiply-accumulate (vector).
  • UMMLA (vector): Unsigned 8-bit integer matrix multiply-accumulate (vector).
  • USMMLA (vector): Unsigned and signed 8-bit integer matrix multiply-accumulate (vector).
  • BFMMLA: BFloat16 floating-point matrix multiply-accumulate into 2x2 matrix.

以Int8类型为例:
image.png

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

image.png
这条指令将第一个源向量中的2x8有符号8位整数矩阵与第二个源向量中的8x2有符号8位整数矩阵相乘。生成的2x2 32位整数矩阵乘积累加到目标向量中的32位整数矩阵累加器中。这相当于对每个目标元素执行8路点积运算。
这条指令在将结果数据类型扩大保持结果精度和128-bit向量寄存器的情况下,执行4x(8次乘+8次加),即32x Int8 MAC,是Dot product指令的2倍,是MLA类指令的四倍。同时,因为单次load到vector寄存器的数据参与的乘加运算更多,提高了运算次数/内存访问的比例。

BF16矩阵乘指令执行的操作为:
image.png

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

image.png

这条对每两个相邻的BFloat16元素对执行两次乘积求和运算,同时将第一个源向量中的2x4 BFloat16矩阵与第二个源向量中的4x2 BFloat16矩阵相乘。中间的求和结果累加到目标向量中的2x2单精度矩阵。这相当于对每个目标元素累加两个2路点积运算。

SVE指令对AI的增强

image.png
SVE指令继承NEON指令中MLA,Dot Product, Matrix Multiply指令,同时提供了支持更大向量寄存器的灵活设计,可以实现单条指令的运算能力。有关SVE的介绍请参见之前的文章。
Arm可伸缩性向量扩展-SVE(上)
Arm可伸缩性向量扩展- SVE(下)
CPU硬件实现的向量长度可以是128-bit的整数倍,SVE指令对MLA,Dot Product, Matrix Multiply的支持,基本是将SVE向量寄存器按每128-bit进行划分,然后每128-bit这些子向量执行与对应NEON指令相同/类似的操作,由此可知,SVE相对NEON,这些MLA, Dot Product, Matrix Multiply指令的操作的MAC数量是以SVE向量寄存器长度除以128-bit的倍数增长的。当然相对NEON,SVE对操作方式和数据类型进一步丰富。

SVE Dot Product支持:

  • SDOT (2-way, indexed): Signed integer indexed dot product.
  • SDOT (2-way, vectors): Signed integer dot product.
  • SDOT (4-way, indexed): Signed integer indexed dot product.
  • SDOT (4-way, vectors): Signed integer dot product.
  • SUDOT: Signed by unsigned integer indexed dot product.
  • UDOT (2-way, indexed): Unsigned integer indexed dot product.
  • UDOT (2-way, vectors): Unsigned integer dot product.
  • UDOT (4-way, indexed): Unsigned integer indexed dot product.
  • UDOT (4-way, vectors): Unsigned integer dot product.
  • USDOT (indexed): Unsigned by signed integer indexed dot product.
  • BFDOT (indexed): BFloat16 floating-point indexed dot product.
  • BFDOT (vectors): BFloat16 floating-point dot product.
  • FDOT (2-way, indexed, FP16 to FP32): Half-precision floating-point indexed dot product.
  • FDOT (2-way, indexed, FP8 to FP16): 8-bit floating-point indexed dot product to half-precision.
  • FDOT (2-way, vectors, FP16 to FP32): Half-precision floating-point dot product.
  • FDOT (2-way, vectors, FP8 to FP16): 8-bit floating-point dot product to half-precision.
  • FDOT (4-way, indexed): 8-bit floating-point indexed dot product to single-precision.
  • FDOT (4-way, vectors): 8-bit floating-point dot product to single-precision

以Int8的Dot Product为例:
image.png

SDOT <Zda>.S, <Zn>.B, <Zm>.B

image.png

这条指令进行有符号整数点积指令计算点积,具体是将第一个源向量中每个32位内的四个有符号8位整数值,与第二个源向量中对应32位元素内的四个有符号8位整数值相乘,然后将扩展后的点积累加到目标向量中对应的32位元素中。
如果CPU硬件实现的SVE向量长度为256-bit, 相对于128-bit NEON,其执行的MAC数量增加一倍。
SVE SDOT指令也支持by element的形式:
image.png

SDOT <Zda>.S, <Zn>.B, <Zm>.B[Index]

这条指令执行有符号整数index点积指令计算点积,将第一个源向量中每个32位元素内的四个有符号8位整数值,与第二个源向量中Index指定的32位元素内的四个有符号8位整数值相乘,然后将扩展后的点积累加到目标向量中对应的32位元素中。
这条指令较清楚地演示了SVE是以128-bit为单位切割的,因为指令中的Index是相对在每128-bit里面的Index。
SVE Dot product还支持Int16类型, 
image.png

UDOT <Zda>.S, <Zn>.H, <Zm>.H

image.png
这条指令进行无符号整数点积指令计算点积,将第一个源向量中每个32位元素内的两个无符号16位整数值与第二个源向量中对应32位元素内的两个无符号16位整数值相乘,然后将扩展后的点积累加到目标向量中对应的32位元素中。

再看一个2路FP8的Dot product, 结果扩大为半精度的指令:
image.png

FDOT <Zda>.S, <Zn>.H, <Zm>.H

image.png

这条指令计算存储在第一个源向量和第二个源向量中每个32位元素的半精度浮点值的fused乘加运算(不进行中间舍入),然后将得到的单精度乘加结果累加到目标向量中对应的单精度元素上。

SVE增加了对单精度,双精度的矩阵乘指令的支持

它执行如下操作(以单精度为例):
image.png
image.png

需要注意的是SVE的Dot Product, Matrix multiply指令是不支持SVE predicate的。

SME Out product

image.png
Armv9-a SME引入了可伸缩矩阵扩展,引入了streaming SME mode,在此模式下,支持更长的向量长度,Out product(外积)指令和用于存储外积结果的二维ZA storage。有关SME请参见相关文章。这里不在赘述。
Arm Scalable Matrix Extension介绍
第二部分: Arm Scalable Matrix Extension (SME)指令

下面只引用其中的几张图片:
2768.Picture5.png

1030.Picture6.png

6545.Picture7.png

SME out product 指令支持predicate:
2072.Picture14.png

以下表格显示了几种数据类型和SVL长度的一条外积并累加或累减指令所做的对应数据类型的MAC(乘累加)数量:
image.png

由此看出SME Out product指令执行的乘加运算的数量是以(SVL/128)的平方倍数增长的。

SME2 Multi-vector

SME2架构引入了对一组Z向量寄存器和ZA数组向量进行操作的多向量指令。它更加有效的利用了SME引入的streaming SVE mode的向量寄存器和ZA storage,以及相应的硬件计算资源。
SME2增加了支持多向量操作数的指令,源操作数和目标操作数可以是以下之一:
• 一组由1个、2个或4个SVE Z向量寄存器组成
• 一组由1个、2个或4个ZA slice组成
• 一组由1个、2个、4个、8个或16个ZA array组成

SME2 Mutli-vector增加了MLA,Dot product运算的支持。以Two ZA single-vectors BF16为例:
image.png

BFDOT ZA.S[<Wv>, <offs>{, VGx2}], { <Zn1>.H-<Zn2>.H }, <Zm>.H

其中每个Op执行前面介绍的SVE Dot product操作,即每个Op执行:
image.png

这条SME2 mutli-vector指令执行相当于两倍于之前SVE dot prodcut指令的操作。
而这条指令的Four ZA single-vectors形式:
BFDOT ZA.S[<Wv>, <offs>{, VGx4}], { <Zn1>.H-<Zn4>.H }, <Zm>.H
执行相当于四倍于之前SVE dot prodcut指令的操作。

除此之外,为了执行2-bit, 4-bit这类低精度的神经网络,SME2增加了LUTIx 查表(lookup table)指令和与之对应的ZT0寄存器,可以实现快速的2-bit, 4-bit的数据转换为8-bit, 16-bit, 32-bit类型,从而进一步运算:
image.png

总结

本文直观详细地介绍了Arm A-profile指令集对日益重要的人工智能应用的支持。可以看出Arm构架向量扩展从最初的MLA指令支持,到Dot product, Matrix Mutliply, 再到SME2 out product, multi-vector指令的支持,提供越来越强大的运算,丰富的数据类型支持,更高的运算/内存访问比例,从而让人工智能应用可以更高效快速灵活地部署在Arm CPU上。下表以输入为Int8, 输出结果为Int32的各指令为例,来展示其乘加数量。
image.png

注:单条指令执行的运算数量不等同于CPU一个cycle可以执行的运算数,其取决于CPU的硬件实现。

推荐阅读
关注数
8647
内容数
61
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息