我被分配任务实现一个优化的矩阵乘法微内核，用C++计算C = A*B，起点是以下代码片段。我遇到了一些反直觉的行为，希望得到一些帮助来更好地理解发生了什么。
void mat_mul(double* A, double* B, double* C) {
for (int n = 0; n < N; ++n) {
for (int k = 0; k < K; ++k) {
for (int m = 0; m < M; ++m) {
C[m + n*M] += A[m + k*M] * B[k + n*K];
}
}
}
}
挑战的条件如下：
- M、N、K在编译时被定义
- K = 128
- M和N可以在编译时选择，但必须小于32
- 矩阵按列主序排列
- 数据对齐可以在编译时决定（默认为32字节对齐）
- 矩阵A的维度为MxK
- 矩阵B的维度为KxN
- 矩阵C的维度为MxN
- 代码必须在AMD EPYC 7742 CPU的单个核心上运行（支持AVX2 + FMA）
- 阻塞是不允许的
通过查看循环的顺序，似乎内存访问顺序已经最小化了缓存未命中，因为我们正在线性迭代缓冲区。
我的第一个想法是尝试向量化代码。这是我得出的结果：
void mat_mul(double* A, double* B, double* C) {
for (int n = 0; n < N; ++n) {
for (int k = 0; k < K; ++k) {
__m256d B_v = _mm256_broadcast_sd(B + k + n*K);
for (int m = 0; m < M; m+=4) {
__m256d A_v = _mm256_load_pd(A + m + k*M);
__m256d C_v = _mm256_load_pd(C + m + n*M);
__m256d rax = _mm256_fmadd_pd(A_v, B_v, C_v);
_mm256_store_pd(C + m + n*M, rax);
}
}
}
}
当M=N=32时，最大性能可达到约23 GFLOPs。当减小N和M时，性能会下降。
经过一段时间的思考，我意识到EPYC 7742的L1d缓存为32KiB，相当于4096个双精度浮点数。理想情况下，我希望所有三个矩阵都能加载到L1缓存中。
这导致了以下优化问题：
最大化N>0，M>0，使得N*M + 128*N + 128 * M ≤ 4096。
我注意到M = 4，N = 8是一个足够好的解决方案，使用2的幂次值。
鉴于M=4，我可以通过保持单个向量作为累加变量来消除m循环。
这个想法产生了以下代码：
void mat_mul(double* A, double* B, double* C) {
__m256d A_v, B_v;
register __m256d C_v;
double* aPtr;
double* bPtr;
double* cPtr;
for (int n = 0; n < N; ++n) {
cPtr = C + n*M;
C_v = _mm256_load_pd(cPtr);
for (int k = 0; k < K; ++k) {
aPtr = A + k*M;
bPtr = B + k + n*K;
B_v = _mm256_broadcast_sd(bPtr);
A_v = _mm256_load_pd(aPtr);
C_v = _mm256_fmadd_pd(A_v, B_v, C_v);
}
_mm256_store_pd(cPtr, C_v);
}
}
我认为这应该会快得多，但是我得到的性能大约为4 GFLOPs，这与使用次优的M=4、N=8运行第一段代码相同。
考虑到在第一种情况下，矩阵太大，无法完全适应L1d高速缓存，这似乎表明代码的第二个版本正在从L2缓存中获取和写入数据，即使矩阵应该适合于L1d高速缓存。
我的怀疑是否正确？如果是这样，这种行为是由于我的思维过程中出现了错误，还是因为我错过了某些原因呢？
请提供一些解释，而不仅仅是发布一个更好的优化代码版本，因为我真的很想更好地理解正在发生的概念。如果您们可以给我一些提示，让我自己查找一些问题，那将是很好的。
谢谢 :)
编辑1：
我尝试遵循@道格和@ElderBug的评论中提出的一些建议。
特别是@doug建议尝试转置B，但是由于矩阵采用列主格式，我找不到实现他们想法的方法。所以我转置A并在一个tmp变量中累加乘积。
这是我想出来的:
void mat_mul(double* A, double* B, double* C) {
__m256d B_v, A_v, rax;
// Transpose A
double AT[K*M];
for(int k=0; k < K; ++k){
for(int m=0; m<M; ++m){
AT[m*K + k] = A[m + k*M];
}
}
// Compute product
for (int n = 0; n < N; ++n) {
for (int m = 0; m < M; ++m) {
double tmp = 0.;
for (int k = 0; k < K; k+=4) {
B_v = _mm256_load_pd(B + n*K + k);
A_v = _mm256_load_pd(AT + m*K + k);
rax = _mm256_mul_pd(A_v, B_v);
double* pv = (double*)&rax;
tmp += pv[0] + pv[1] + pv[2] + pv[3];
}
C[n*M + m] = tmp;
}
}
}
在M=4，N=8的情况下，此代码仍然以约4 GFLOPs的速度运行。问题在于需要减少rax向量。我没有找到更有效的方法来实现。
如果我删除tmp += pv[0] + pv[1] + pv[2] + pv[3];，性能将翻倍，但是当M=N=32时，我只能达到14 GFLOPs的峰值，因此这仍然比我的天真向量化方法差。
如果有人有其他建议或评论，将非常感激。
编辑2：
我忘记提到我正在使用以下优化标志编译代码：icc (ICC) 2021.5.0 20211109
-O3 -funroll-loops -std=c++11 -mavx2 -march=native -fma -ftz -fomit-frame-pointer
编辑3：
目标是以最佳方式实现这个串行微内核，以便我可以将其用于阻塞并行矩阵乘法。根据我的计算，理论峰值应该为46 GFLOPS，因此我得到了大约50%的性能。OpenBLAS获得了约40 GFLOPs，所以32-35左右会很好。如果有人能解释一下为什么我尝试的某些选项不起作用，那就太好了，这样我就可以考虑如何修复它们。
编辑4：
感谢@ElderBug的评论，我对我的第一个向量化实现中的存储操作管理进行了一些巧妙的展开，并成功达到了40GFLOps，这与OpenBLAS相当！
读一下这篇文章也有帮助：https://github.com/flame/blis/blob/master/docs/KernelsHowTo.md#gemm-microkernel
void mat_mul(double* __restrict__ A, double* __restrict__ B, double* C)的不同版本。__restrict__关键字告诉编译器 A、B 和 C 是单独的矩阵，从而实现更多优化。这如何改变您的系统上的结果？ - John Zwinck