前文讲了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（点积）指令的操作
- By vector: 放在一个向量寄存器里的4x4 Int8矩阵a和放在另一个向量寄存器里的4x4 Int8矩阵b，进行矩阵a的行（n）与矩阵b的列（n，列号与行号相同）的点积运算，得到4个32-bit的结果。或是
- By element: 矩阵a的其中一行与矩阵b的每列的点积运算，得到4个32-bit的结果。
对于SVE2 Int8 Dot product指令，相等于执行(VL / 128)个以上操作。
例如SVE2向量长度VL=256-bit，它执行：
- 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结果。或是
- By element: Zm中的矩阵a1的其中一行与矩阵b1的每列的点积运算，得到4个32-bit的结果；Zn中的矩阵a2的其中一行（与a1的行号相同）与矩阵b2的每列的点击运算，得到4个32-bit的结果。总共得到8个32-bit结果。
更多有关点积指令，请看Arm构架如何让AI应用高效运行于CPU (1)
以SVE2 Int8 Dot Product为例，假设矩阵的行数和列数可以与向量长度VL匹配（不需要做tail处理），A矩阵和B矩阵的矩阵乘（得到C矩阵）可以表达为:
其中标识出来的部分的功能可以使用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矩阵以这样的方式继续分块，为了简略，图中没有画出所有的块。
B矩阵可以处理为：
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矩阵以这样的方式继续分块，为了简略，图中没有画出所有的块。
点积运算
接下来进行运算：
将经重排过的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]），以此类推。
将经重排过的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]），以此类推。
然后使用SVE2 Dot Product by element指令，
UDOT <Zda>.S, <Zn>.B, <Zm>.B[0]来执行如下图所示的操作：
A块的0（0）与B块0做Dot Product by element运算，这会计算出8个32-bit的点积结果（VL=256-bit），这些结果是如下图所示的C矩阵的这些元素的部分结果（中间结果，不是最终结果）。
- 再利用1，2，3类似的方式计算：
A块的0（1）与B块0做Dot Product by element运算，这会计算出8个32-bit的点积结果（VL=256-bit），这些结果是如下图所示的C矩阵的这些元素的部分结果（中间结果，不是最终结果）。
重复这样的过程，计算：
• A块的0（2）与B块0做Dot Product by element运算
• A块的0（3）与B块0做Dot Product by element运算
经过以上步骤，会计算出C矩阵以下元素的中间结果：
- 计算：
• 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矩阵以下元素的中间结果：
- 计算：
• A块的2（0）与B块3做Dot Product by element运算（此过程会累加上面运算的C中间结果）
这会将C矩阵中的以下元素更新为：
• 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矩阵的元素更新为：
- 继续在K维度对A矩阵的块与B矩阵的块进行类似的迭代运算，完成K维度迭代之后，可以得到C矩阵的以下元素的最终结果：
- 以同样的方式计算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>
ZIP2 <Zd>.<T>, <Zn>.<T>, <Zm>.<T>
其中T的类型可以是8,16,32,64或是128 bit, 即交织的粒度可以是byte, half word,word, double word..
举个例子，如果要将以下矩阵A
重排为利用dot product矩阵运算所需要的8块4x4的矩阵
重排后在内存中存放的视图为：
重排后对A矩阵的4x4 block访问在内存里是连续的，这样可以方便load到向量寄存器中，而且对内存，cache更加友好。
这个过程是：
使用
LD1B { <Zt>.B }, <Pg>/Z, [<Xn|SP>{, #<imm>, MUL VL}]指令将原来矩阵的四行中的VL长度数据load到4个Z寄存器中：
利用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为
利用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中得到我们想要的重排数据分布。
- 利用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指令加速矩阵乘运算。
