11

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

如何使用Arm 向量指令加速矩阵乘 (1) - SVE2 Dot Product

前文讲了Arm用于加速AI, ML应用的向量和SME指令,本文介绍如何使用这些指令来实现矩阵乘。
A(M x K)矩阵和B (K x N)矩阵的矩阵乘(得到C矩阵 M x N)可以表达为:

for i in range (M) :
      for j in range (N) :
            for k in range (K) :
                  C[i][j] += C[i][j]+A[i][k]*B[k][j]

Int8的Dot product(点积)指令的操作

  1. By vector: 放在一个向量寄存器里的4x4 Int8矩阵a和放在另一个向量寄存器里的4x4 Int8矩阵b,进行矩阵a的行(n)与矩阵b的列(n,列号与行号相同)的点积运算,得到4个32-bit的结果。或是
  2. By element: 矩阵a的其中一行与矩阵b的每列的点积运算,得到4个32-bit的结果。

对于SVE2 Int8 Dot product指令,相等于执行(VL / 128)个以上操作。
例如SVE2向量长度VL=256-bit,它执行:

  1. By vector: 放在一个向量寄存器Zm的低128-bit里的4x4 Int8矩阵a1和放在另一个向量寄存器Zn里的4x4 Int8矩阵b1,进行矩阵a1的行(n)与矩阵b1的列(n)的点积运算,得到4个32-bit的结果;向量寄存器Zm高128-bit里的4x4 Int8矩阵a2和放在另一个向量寄存器Zn里的4x4 Int8矩阵b2,进行矩阵a2的行(n)与矩阵b2的列(n)的点积运算,得到4个32-bit的结果。总共得到8个32-bit结果。或是
  2. By element: Zm中的矩阵a1的其中一行与矩阵b1的每列的点积运算,得到4个32-bit的结果;Zn中的矩阵a2的其中一行(与a1的行号相同)与矩阵b2的每列的点击运算,得到4个32-bit的结果。总共得到8个32-bit结果。
    image.png
    image.png

更多有关点积指令,请看Arm构架如何让AI应用高效运行于CPU (1)

以SVE2 Int8 Dot Product为例,假设矩阵的行数和列数可以与向量长度VL匹配(不需要做tail处理),A矩阵和B矩阵的矩阵乘(得到C矩阵)可以表达为:
image.png

其中标识出来的部分的功能可以使用SVE2的Dot Product的by element来实现。为了优化内存访问和数据重用,我们优化为进行back to back的Dot Product, 如果CPU的硬件pipeline实现支持同时多条Dot Product指令的执行,这样可以更好的利用这些硬件资源。
我们以进行3个vector长度 back to back的SVE2 Int8 Dot Product,VL=256-bit 为例,其过程可以图示为:

A矩阵和B矩阵的分割和重排

首先,需要对A矩阵和B矩阵进行分割处理(为了内存访问更友好,需要对A,B矩阵在内存中进行重排):
A矩阵可以处理为:

A矩阵块0由[a0_0, a0_1 , a0_2, a0_3, a1_0, a1_1 , a1_2, a1_3, a2_0, a2_1 , a2_2, a2_3, a3_0, a3_1 , a3_2, a3_3]这些4x4 byte (128-bit)元素组成。块1,2,3以此类推。A矩阵以这样的方式继续分块,为了简略,图中没有画出所有的块。
gemm-sve_interleaved_u8u32_3VLx8_1.jpg

B矩阵可以处理为:
gemm-sve_interleaved_u8u32_3VLx8_2.jpg

B矩阵块0由[b0_0, b1_0, b2_0, b3_0, b0_1, b1_1, b2_1, b3_1, b0_2, b1_2, b2_2, b3_2, b0_3, b1_3, b2_3, b3_3, b0_4, b1_4, b2_4, b3_4, b0_5, b1_5, b2_5, b3_5, b0_6, b1_6, b2_6, b3_6, b0_7, b1_7, b2_7, b3_7] 这些元素(总共为VL bits)组成。每3块(3xVL,如块0,1,2)为一组。B矩阵以这样的方式继续分块,为了简略,图中没有画出所有的块。

点积运算

接下来进行运算:

  1. 将经重排过的A矩阵的块0 load到一个SVE2向量寄存器,(经重排之后的A矩阵的块在内存中会地址连续存储),可以使用LD1RQB指令

    LD1RQB { <Zt>.B }, <Pg>/Z, [<Xn|SP>, <Xm>]

    这条指令从内存中load 16 byte(128-bit)到一个SVE2向量寄存器的每一个128-bit段(重复这16 byte),如下图中vector中红色部分,0(0)表示A 矩阵块0的第0个4 byte(即 [a0_0, a0_1 , a0_2, a0_3]), 0(1)表示A 矩阵块0的第1个4 byte(即 [a1_0, a1_1 , a1_2, a1_3]),以此类推。

  2. 将经重排过的B矩阵的块0 load到一个SVE2向量寄存器,(经重排之后的A矩阵的块在内存中会地址连续存储),可以使用LD1B指令

    LD1B { <Zt>.B }, <Pg>/Z, [<Xn|SP>{, #<imm>, MUL VL}]

    这条指令从内存中load VL byte到一个SVE2向量寄存器,如下图中vector中蓝色部分,0(0)表示B 矩阵块0的第0个4 byte(即 [b0_0, b1_0, b2_0, b3_0]), 0(1)表示B 矩阵块0的第1个4 byte(即 [b0_1, b1_1, b2_1, b3_1]),以此类推。

  3. 然后使用SVE2 Dot Product by element指令,

    UDOT <Zda>.S, <Zn>.B, <Zm>.B[0]

    来执行如下图所示的操作:
    gemm-sve_interleaved_u8u32_3VLx8_3.jpg

A块的0(0)与B块0做Dot Product by element运算,这会计算出8个32-bit的点积结果(VL=256-bit),这些结果是如下图所示的C矩阵的这些元素的部分结果(中间结果,不是最终结果)。
gemm-sve_interleaved_u8u32_3VLx8_4.jpg

  1. 再利用1,2,3类似的方式计算:
    gemm-sve_interleaved_u8u32_3VLx8_5.jpg
    A块的0(1)与B块0做Dot Product by element运算,这会计算出8个32-bit的点积结果(VL=256-bit),这些结果是如下图所示的C矩阵的这些元素的部分结果(中间结果,不是最终结果)。
    gemm-sve_interleaved_u8u32_3VLx8_6.jpg

重复这样的过程,计算:
• A块的0(2)与B块0做Dot Product by element运算
• A块的0(3)与B块0做Dot Product by element运算

经过以上步骤,会计算出C矩阵以下元素的中间结果:
gemm-sve_interleaved_u8u32_3VLx8_7.jpg

  1. 计算:
    • A块的0(0)与B块1做Dot Product by element运算
    • A块的0(1)与B块1做Dot Product by element运算
    • A块的0(2)与B块1做Dot Product by element运算
    • A块的0(3)与B块1做Dot Product by element运算
    • A块的0(0)与B块2做Dot Product by element运算
    • A块的0(1)与B块2做Dot Product by element运算
    • A块的0(2)与B块2做Dot Product by element运算
    • A块的0(3)与B块2做Dot Product by element运算

经过以上步骤,会计算出C矩阵以下元素的中间结果:
gemm-sve_interleaved_u8u32_3VLx8_8.jpg

  1. 计算:
    • A块的2(0)与B块3做Dot Product by element运算(此过程会累加上面运算的C中间结果)
    gemm-sve_interleaved_u8u32_3VLx8_9.jpg

这会将C矩阵中的以下元素更新为:
gemm-sve_interleaved_u8u32_3VLx8_10.jpg

• A块的2(1)与B块3做Dot Product by element运算(此过程会累加上面运算的C中间结果)
• A块的2(2)与B块3做Dot Product by element运算(此过程会累加上面运算的C中间结果)
• A块的2(3)与B块3做Dot Product by element运算(此过程会累加上面运算的C中间结果)

• A块的2(0)与B块4做Dot Product by element运算(此过程会累加上面运算的C中间结果)
• A块的2(1)与B块4做Dot Product by element运算(此过程会累加上面运算的C中间结果)
• A块的2(2)与B块4做Dot Product by element运算(此过程会累加上面运算的C中间结果)
• A块的2(3)与B块4做Dot Product by element运算(此过程会累加上面运算的C中间结果)

• A块的2(0)与B块5做Dot Product by element运算(此过程会累加上面运算的C中间结果)
• A块的2(1)与B块5做Dot Product by element运算(此过程会累加上面运算的C中间结果)
• A块的2(2)与B块5做Dot Product by element运算(此过程会累加上面运算的C中间结果)
• A块的2(3)与B块5做Dot Product by element运算(此过程会累加上面运算的C中间结果)

经上面步骤,C矩阵的元素更新为:
gemm-sve_interleaved_u8u32_3VLx8_11.jpg

  1. 继续在K维度对A矩阵的块与B矩阵的块进行类似的迭代运算,完成K维度迭代之后,可以得到C矩阵的以下元素的最终结果:
    gemm-sve_interleaved_u8u32_3VLx8_12.jpg
  2. 以同样的方式计算C矩阵的其他元素。

代码实现如下:

 "whilelt p0.b, XZR, x25\n"
  "ld1b { z21.b }, p2/Z, [x28]\n"
  "ld1b { z26.b }, p2/Z, [x28, #1, MUL VL]\n"
  "ld1b { z25.b }, p2/Z, [x28, #2, MUL VL]\n"
  "ld1b { z24.b }, p2/Z, [x28, #3, MUL VL]\n"
  "ld1b { z20.b }, p2/Z, [x28, #4, MUL VL]\n"
  "ld1b { z23.b }, p2/Z, [x28, #5, MUL VL]\n"
  "ld1rqb { z0.b }, p0/Z, [x24]\n"
  "ld1b { z22.b }, p2/Z, [x28, #6, MUL VL]\n"
  "add x24, x24, #0x10\n"
  "udot z16.s, z21.b, z0.b[0]\n"
  "ld1b { z21.b }, p2/Z, [x28, #7, MUL VL]\n"
  "addvl x28, x28, #16\n"
  "udot z17.s, z26.b, z0.b[0]\n"
  "udot z18.s, z25.b, z0.b[0]\n"
  "udot z19.s, z24.b, z0.b[0]\n"
  "udot z16.s, z20.b, z0.b[1]\n"
  "ld1b { z20.b }, p2/Z, [x28, #-8, MUL VL]\n"
  "ld1b { z26.b }, p2/Z, [x28, #-7, MUL VL]\n"
  "ld1b { z25.b }, p2/Z, [x28, #-6, MUL VL]\n"
  "ld1b { z24.b }, p2/Z, [x28, #-5, MUL VL]\n"
  "udot z17.s, z23.b, z0.b[1]\n"
  "ld1b { z23.b }, p2/Z, [x28, #-4, MUL VL]\n"
  "udot z18.s, z22.b, z0.b[1]\n"
  "ld1b { z22.b }, p2/Z, [x28, #-3, MUL VL]\n"
  "udot z19.s, z21.b, z0.b[1]\n"
  "ld1b { z21.b }, p2/Z, [x28, #-2, MUL VL]\n"
  "udot z16.s, z20.b, z0.b[2]\n"
  "ld1b { z20.b }, p2/Z, [x28, #-1, MUL VL]\n"
  "udot z17.s, z26.b, z0.b[2]\n"
  "udot z18.s, z25.b, z0.b[2]\n"
  "udot z19.s, z24.b, z0.b[2]\n"
  "udot z16.s, z23.b, z0.b[3]\n"
  "udot z17.s, z22.b, z0.b[3]\n"
  "udot z18.s, z21.b, z0.b[3]\n"
  "udot z19.s, z20.b, z0.b[3]\n"
  "tbnz %x[flags], #31, 8f\n"
  "udot z11.s, z0.b, z15.b\n"
  "8:"

以上汇编代码高效利用了SVE2提供的Z寄存器数量,穿插SVE2内存访问指令和Back to Back Dot product 运算指令,以便有效利用CPU硬件pipeline。

可以参见Arm compute library和KleidiAI中的代码:

https://github.com/ARM-softwa...

https://gitlab.arm.com/kleidi...
KleidiAI中的实现来处理了量化的一些操作,会和介绍的实现稍有差别。
Arm Compute Library还包含了比较丰富的量化,较小的K值处理, NEON,等不同的kernel实现:
https://github.com/ARM-softwa...
https://github.com/ARM-softwa...
但基本的使用Dot Product by element的方式基本不变。

如何快速重排矩阵?

前面讲到需要对A矩阵和B矩阵进行重排,那么如何可以快速地将矩阵重排?
答案是可以借助arm向量指令的ZIP1和ZIP2指令。ZIP指令是用来交织两个向量的:

ZIP1 <Zd>.<T>, <Zn>.<T>, <Zm>.<T>

gemm-sve_interleaved_u8u32_3VLx8_13.jpg

ZIP2 <Zd>.<T>, <Zn>.<T>, <Zm>.<T>

gemm-sve_interleaved_u8u32_3VLx8_14.jpg

其中T的类型可以是8,16,32,64或是128 bit, 即交织的粒度可以是byte, half word,word, double word..

举个例子,如果要将以下矩阵A
intervleave_0.JPG

重排为利用dot product矩阵运算所需要的8块4x4的矩阵
intervleave_2.jpg
重排后在内存中存放的视图为:
intervleave_3.jpg

重排后对A矩阵的4x4 block访问在内存里是连续的,这样可以方便load到向量寄存器中,而且对内存,cache更加友好。
这个过程是:

  1. 使用

    LD1B { <Zt>.B }, <Pg>/Z, [<Xn|SP>{, #<imm>, MUL VL}]

    指令将原来矩阵的四行中的VL长度数据load到4个Z寄存器中:
    intervleave_4.jpg

  2. 利用ZIP1,ZIP2和T为32-bit(word)的方式交织这些Z寄存器

    ZIP1 Ze.S, Za.S, Zb.S
    ZIP2 Zf.S,  Za.S, Zb.S
    ZIP1 Zg.S, Zc.S,  Zd.S
    ZIP2 Zh.S, Zc.S, Zd.S

    经此过程得到Ze, Zf, Zg, Zh为
    intervleave_5.jpg

  3. 利用ZIP1,ZIP2和T为64-bit(double word)的方式交织这些Z寄存器

    ZIP1 Za.D, Ze.D, Zg.D
    ZIP2 Zb.D,  Ze.D, Zg.D
    ZIP1 Zc.D, Zf.D,  Zh.D
    ZIP2 Zd.D, Zf.D, Zh.D

    经此过程在Za,Zb,Zc,Zd中得到我们想要的重排数据分布。
    intervleave_6.jpg

  4. 利用ST1W将Z寄存器中的值存入重排矩阵内存中。

在Arm compute library中类似的interleave处理:
https://github.com/ARM-softwa...

结语

利用SVE2的ZIP交织指令可以快速重排矩阵,Dot Product by element指令可以快速迭代计算重排之后的矩阵乘结果,显著提升CPU进行Int8矩阵乘运算效率和速度。对于其他数据类型的Dot product也可以采取类似的处理。

后续文章介绍如何利用SVE2矩阵乘向量指令和SME2 outer product指令加速矩阵乘运算。

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