使用 AVX 的平铺矩阵乘法

Question

我编写了以下 C 函数，用于使用 tiling/blocking 和 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 }

Answer 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 个不同的列（B[k][j+0..3] 通过来自 B[k][j] 的向量加载），并从 A[i][k] 广播 1 个双倍。请记住，您需要并行 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 嵌套循环。（不过只是标量，未矢量化）。

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

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

其他 matmul 策略包括在内循环内添加目标，因此您在内循环内加载 C 和 A，并提升来自 B 的负载。（或者 B 和 A 交换了，我忘记了。）What Every Programmer Should Know About Memory? has a vectorized cache-blocked example that works this way in an appendix, for SSE2 __m128d vectors. https://www.akkadia.org/drepper/cpumemory.pdf

使用 AVX 的平铺矩阵乘法

Tiled Matrix Multiplication using AVX

c

performance

simd

avx

matrix-multiplication