使用 AVX 的平铺矩阵乘法

Posted 2023-02-16

【中文标题】使用 AVX 的平铺矩阵乘法

【英文标题】：Tiled Matrix Multiplication using AVX

【发布时间】：2019-11-23 16:15:44

【问题描述】：

我编写了以下 C 函数，用于使用平铺/分块和 AVX 向量将两个 NxN 矩阵相乘以加快计算速度。现在，当我尝试将 AVX 内在函数与平铺结合时，我遇到了分段错误。知道为什么会这样吗？

另外，矩阵 B 是否有更好的内存访问模式？也许先转置它，甚至改变 k 和 j 循环？因为现在，我正在逐列遍历它，这在空间局部性和缓存行方面可能不是很有效。

  1 void mmult(double A[SIZE_M][SIZE_N], double B[SIZE_N][SIZE_K], double C[SIZE_M][SIZE_K])
  2 
  3   int i, j, k, i0, j0, k0;
  4   // double sum;
  5   __m256d sum;
  6   for(i0 = 0; i0 < SIZE_M; i0 += BLOCKSIZE) 
  7   for(k0 = 0; k0 < SIZE_N; k0 += BLOCKSIZE) 
  8   for(j0 = 0; j0 < SIZE_K; j0 += BLOCKSIZE) 
  9       for (i = i0; i < MIN(i0+BLOCKSIZE, SIZE_M); i++) 
 10         for (j = j0; j < MIN(j0+BLOCKSIZE, SIZE_K); j++) 
 11           // sum = C[i][j];
 12           sum = _mm256_load_pd(&C[i][j]);
 13           for (k = k0; k < MIN(k0+BLOCKSIZE, SIZE_N); k++) 
 14             // sum += A[i][k] * B[k][j];
 15             sum = _mm256_add_pd(sum, _mm256_mul_pd(_mm256_load_pd(&A[i][k]), _mm256_broadcast_sd(&B[k][j])));
 16           
 17           // C[i][j] = sum;
 18           _mm256_store_pd(&C[i][j], sum);
 19         
 20       
 21   
 22   
 23   
 24 

【问题讨论】：

@zx485：嗯？注意循环边界：内部循环在整个矩阵的 BLOCKSIZE 切片上，允许在 L1d 缓存中仍然很热的情况下重用数据。外部循环使用+=BLOCKSIZE。但是，IIRC 应该只有 5 个嵌套循环，而不是 6 个，而且细节可能是错误的。

@zx485 好吧，BLOCKSIZE 有一个固定的（小）值，比如 8 或 16，而 N 可能要大得多。所以，它并不是真正的 O(N^6)，但它应该是“通常的”O(N^3)。

【参考方案1】：

_mm256_load_pd 是需要对齐的加载，但在加载 4 个双精度值的 32 字节向量的最内层循环中，您只是通过 k++ 而不是 k+=4 步进。所以它会出现故障，因为每 4 个负载中有 3 个未对齐。

你不想做重叠加载，你真正的错误是索引；如果您的输入指针是 32 字节对齐的，您应该能够继续使用 _mm256_load_pd 而不是 _mm256_loadu_pd。因此，使用 _mm256_load_pd 成功捕获了您的错误，而不是正常工作，但给出了数字错误的结果。

---

您对四个 row*column 点积进行矢量化的策略（以生成 C[i][j+0..3] 向量）应该从 4 个不同的列加载 4 个连续的双精度数（B[k][j+0..3] 通过来自 B[k][j] 的向量加载)，并从A[i][k] 广播 1 个 double。请记住，您需要并行 4 个点积。

另一种策略可能涉及到最后的水平总和到标量C[i][j] += horizontal_add(__m256d)，但我认为这需要先转置一个输入，以便行向量和列向量都在一个点积的连续内存中。但是，您需要在每个内部循环结束时对水平总和进行随机播放。

您可能还想使用至少 2 个sum 变量，这样您就可以一次读取整个缓存行，并隐藏内部循环中的 FMA 延迟，并有望成为吞吐量瓶颈。 或者更好地并行执行 4 或 8 个向量。 所以你生成 C[i][j+0..15] 为 sum0、sum1、sum2、sum3。 （或者使用 __m256d 的数组；编译器通常会完全展开一个 8 的循环并将数组优化到寄存器中。）

---

我认为您只需要 5 个嵌套循环来阻止行和列。尽管显然 6 个嵌套循环是一个有效的选项：请参阅loop tiling/blocking for large dense matrix multiplication，它在问题中有一个 5 个嵌套循环，但在答案中有一个 6 个嵌套循环。 （只是标量，而不是矢量化）。

这里除了行*列点积策略可能还有其他bug，我不确定。

---

如果您使用的是 AVX，则可能还需要使用 FMA，除非您需要在 Sandbybridge/Ivybridge 和 AMD Bulldozer 上运行。 （打桩机及以后有 FMA3）。

其他 matmul 策略包括添加到内部循环内的目标中，以便您在内部循环内加载 C 和 A，并提升来自 B 的负载。 （或者 B 和 A 交换了，我忘记了。）What Every Programmer Should Know About Memory? 有一个矢量化缓存阻塞示例，在附录中以这种方式工作，用于 SSE2 __m128d 向量。 https://www.akkadia.org/drepper/cpumemory.pdf