riscv平台优化矩阵乘(基于blislab优化实践)
riscv平台优化矩阵乘(基于blislab优化实践)
本文基于blislab与openblas项目，讲述在riscv平台上优化矩阵乘优化过程。
为了方便riscv-rvv加速，将blislab项目中的double数据类型改为float（即测试单精度浮点性能），且例子都是列主序的。
1 blislab项目
Blislab是一个开源教学项目，教你一步一步优化矩阵乘。项目地址：https://github.com/flame/blislab
在blislab项目的基础上，我进行了一些删减优化（删成最简的代码，只支持x86或riscv的linux OS），这样项目更清晰明了，方便入门。
项目地址：
https://github.com/surez-ok/blislab_riscv
上手请参照第3节开始上手
2 优化步骤
step0：循环交换:
循环交换即交换嵌套循序，循环交换的主要目的是：对于多维数组的元素尽量能顺序访问。这样可以改善内存访问的空间局部性，对缓存更加友好，这种转换有助于消除内存带宽和内存时延瓶颈。
ref版本（基础版本）如下：
void bl_sgemm_ref(
int m,
int n,
int k,
float *XA,
int lda,
float *XB,
int ldb,
float *XC,
int ldc
)
{
// Local variables.
int i, j, p;
float alpha = 1.0, beta = 1.0;
// Sanity check for early return.
if ( m == 0 || n == 0 || k == 0 ) return;
// Reference GEMM implementation.
// ijp
for ( i = 0; i < m; i ++ ) {
for ( j = 0; j < n; j ++ ) {
for ( p = 0; p < k; p ++ ) {
XC[ j * ldc + i ] += XA[ p * lda + i ] * XB[ j * ldb + p ];
}
}
}
}
调整i、j、p的顺序，对结果没有影响，但是可以影响性能。
以下测试不同循环顺序的性能（m, n, k取256）
| 顺序 | 某款riscv芯片 |
|---|---|
| ijp | 25.63 MFLOPS |
| ipj | 14.89 MFLOPS |
| jip | 26.22 MFLOPS |
| jpi | 159.87 MFLOPS |
| pij | 14.92 MFLOPS |
| pji | 147.96 MFLOPS |
这样得到性能最高的循环顺序jpi，性能最差的循序为ipj
原因分析：
性能最好的jpi模式的代码如下：
void bl_sgemm(
int m,
int n,
int k,
float *A,
int lda,
float *B,
int ldb,
float *C, // must be aligned
int ldc // ldc must also be aligned
)
{
int i, j, p;
// Early return if possible
if ( m == 0 || n == 0 || k == 0 ) {
printf( "bl_sgemm(): early return\n" );
return;
}
// jpi
for ( j = 0; j < n; j ++ ) { // Start 2-nd loop
for ( p = 0; p < k; p ++ ) { // Start 1-st loop
for ( i = 0; i < m; i ++ ) { // Start 0-th loop
//C[ j * ldc + i ] += A[ p * lda + i ] * B[ j * ldb + p ];
C( i, j ) += A( i, p ) * B( p, j ); //Each operand is a MACRO defined in bl_sgemm() function.
} // End 0-th loop
} // End 1-st loop
} // End 2-nd loop
}
考虑内存循环jpi的计算过程，如下图：
很显然，jpi顺序适合cache预取，这种方式的一级缓存命中率高于ipj模式。
测试ipj顺序在不同数据规模下的表现，从图中可以看出，随着数据规模增大，整体性能在下降（中间的抖动可能与cache换出换入有关）
可以看到：随着数据规模增加，flops将会衰减（测试ipj模式是因为这种模式会较快的达到cache限制，方便测试，jpi性能图如下图类似，只是达到性能下降较慢）
为什么矩阵变大后flops会快速衰减呢？
当A、B 矩阵小于 L2 cache 时，gemm只需要从 DDR中读取 A、B 大小的内存；但是当 AB 大于 L2 cache 时，由于行主序的 B 或者列主序的 A 不是内存连续的，gemm从 DDR读取的内存数超过 A、B 的大小，造成cache miss增加，于是性能恶化。
step1：基本的分块:
解决step0中的问题，一个解决思路是：将大矩阵拆分为小矩阵，这样小矩阵可以完全放到cache中。
基本分块：
先将矩阵分块（m n方向进行分块，k方向不分块），然后展开：
如下图所示：
A矩阵沿着m方向分块，每个小矩阵为 （MR， k）
B矩阵沿着n方向分块，每个小矩阵为 （ k， NR）
小矩阵（MR， k）与 小矩阵 （ k， NR）相乘得到C中的小块矩阵（MR，NR），如小黄
依次类推...
写成如下方式：
void AddDot( int k, float *A, int lda, float *B, int ldb, float *result ) {
int p;
for ( p = 0; p < k; p++ ) {
*result += A( 0, p ) * B( p, 0 );
}
}
void AddDot_MRxNR( int k, float *A, int lda, float *B, int ldb, float *C, int ldc )
{
int ir, jr;
int p;
for ( jr = 0; jr < DGEMM_NR; jr++ ) {
for ( ir = 0; ir < DGEMM_MR; ir++ ) {
AddDot( k, &A( ir, 0 ), lda, &B( 0, jr ), ldb, &C( ir, jr ) );
}
}
}
void bl_sgemm(
int m,
int n,
int k,
float *A,
int lda,
float *B,
int ldb,
float *C, // must be aligned
int ldc // ldc must also be aligned
)
{
int i, j, p;
int ir, jr;
// Early return if possible
if ( m == 0 || n == 0 || k == 0 ) {
printf( "bl_sgemm(): early return\n" );
return;
}
for ( j = 0; j < n; j += DGEMM_NR ) { // Start 2-nd loop
for ( i = 0; i < m; i += DGEMM_MR ) { // Start 1-st loop
AddDot_MRxNR( k, &A( i, 0 ), lda, &B( 0, j ), ldb, &C( i, j ), ldc );
} // End 1-st loop
} // End 2-nd loop
}
其中AddDot_MRxNR函数完成(MR, k)与 (k, NR)矩阵的乘法。
可以测试MR NR不同的取值的性能，可见：仅仅进行简单分块并不能提高性能，因为其循环层增加了，并且计算时是等价的，需要进一步优化。
比如：mr = 2, nr = 2的代码：
使用一些优化技巧：
- 循环展开
- 寄存器缓存（寄存器结果重复利用）
mr = 2 nr = 2 （unrool版本）代码如下：
void AddDot_2x2( int k, float *A, int lda, float *B, int ldb, float *C, int ldc )
{
register float C00 = 0.0f, C01= 0.0f, C10= 0.0f, C11 = 0.0f;
int p;
for (p = 0; p < k; p++) {
C00 += A(0, p) * B(p, 0);
C01 += A(0, p) * B(p, 1);
C10 += A(1, p) * B(p, 0);
C11 += A(1, p) * B(p, 1);
}
C(0, 0) += C00;
C(0, 1) += C01;
C(1, 0) += C10;
C(1, 1) += C11;
}
void bl_sgemm(
int m,
int n,
int k,
float *A,
int lda,
float *B,
int ldb,
float *C, // must be aligned
int ldc // ldc must also be aligned
)
{
int i, j, p;
int ir, jr;
// Early return if possible
if ( m == 0 || n == 0 || k == 0 ) {
printf( "bl_sgemm(): early return\n" );
return;
}
for ( j = 0; j < n; j += 2 ) { // Start 2-nd loop
for ( i = 0; i < m; i += 2 ) { // Start 1-st loop
AddDot_2x2( k, &A( i, 0 ), lda, &B( 0, j ), ldb, &C( i, j ), ldc );
} // End 1-st loop
} // End 2-nd loop
}
测试结果如下：
可见：当矩阵规模较少（k较小）时，unrool版本比ref版本有1x以上提升，但是当矩阵超过一定规模后，性能避免不了的还是下降了，这是因为当k比较大时，矩阵分成的小矩阵在cache里放不下了，所以需要更优的分块方式。
step2：cache级别的分块：
step2相较于step1的区别是：在k方向上也进行了分块。
解决思路依旧是：
将大矩阵拆分为小矩阵，保证小矩阵能放在cache里，这样可以提高访问性能。
下图示例表示分块原理：
所以，可以将大矩阵拆为小矩阵，确保分割后的小矩阵可以全部放到cache中，然后访问。
其伪代码如下：
for jc = 0; jc < n; step = NC
for ic = 0; ic < m; step = MC
for pc = 0; pc < k; step = KC
pack_A(ic, pc) // 准备好pack_A矩阵
pack_B(pc, jc) // 准备好pack_B矩阵
Kernel_op(A(ic, pc), (pc, jc)) // 计算pack_A * pack_B
endfor
endfor
endfor
上述操作可以进行优化（尽量减少内存循环计算量）。
for jc = 0; jc < n; step = NC // loop5
for pc = 0; pc < k; step = KC // loop4
pack_B(jc, pc)
for ic = 0; ic < m; step = MC // loop3
pack_A(ic, pc)
Kernel_op(A(ic, pc), (pc, jc)) // macro-kernel Loop2-loop0
endfor // loop3
endfor // loop4
endfor // loop5
其流程如下(来自openblas流程图)：
以上参数需要满足一定的限制：
MC * KC 小于L2的1/2
实际测试：
从下图测试数据可以看到（见橙色部分），基于step2进行分块后，随着数据规模增大，flops可以保持稳定；
step3：kernel使用SIMD指令优化：
对于step2中的micro kernel部分，可以使用riscv rvv指令优化kernel乘法。
测试结果如下，可见与step2相比，使用rvv指令优化有巨大的提升：
3 开始上手
以step0为例（其它几个例子也是相同操作）：
$ cd step0
# 编译，默认是intel平台，gcc编译
# 如果需要编译riscv平台，需要修改Makefile，将Makefile的前几行中 CROSS_COMPILE 配置为正确的交叉编译器
$ make
# 测试
$ bash run_bl_sgemm.sh
