C++优化矩阵乘法微内核L1缓存使用

Question

我的任务是实现一个优化的矩阵乘法微内核，它从以下代码片段开始用 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个double。理想情况下，我希望所有三个矩阵都加载到 L1 缓存中。

这会产生以下优化问题：

最大化 N>0,M>0 使得 N*M + 128*N + 128 * M ≤ 4096.

我注意到 M = 4，N = 8 是一个足够好的解决方案，它使用二次幂值。

鉴于 M=4，我可以通过保持单个向量作为累积变量来消除 m-loop。

这个想法产生了以下代码：

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：

我尝试遵循@doug 和@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;
    }
  }
}

这仍然以大约 4 GFLOPs 运行，M=4，N=8。房间里的大象似乎在减少 rax 向量。我无法找到更有效地做到这一点的不同方法。

如果我删除

tmp += pv[0] + pv[1] + pv[2] + pv[3];

性能翻倍，但我达到了 14 GFLOPs 的峰值，M=N=32，所以这仍然比我天真的矢量化方法差。

如果有人有任何进一步的建议/意见，他们将不胜感激。

编辑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 GFLOP，所以 32-35ish 会很好。

如果有人可以提供一些见解，说明为什么我尝试过的某些选项不起作用，以便我可以考虑如何修复它们，那就太好了。