使用 AVX 的平铺矩阵乘法

Question

我编写了以下 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 }

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] ）加载 4 个连续的双打通过来自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 个嵌套循环是一个有效选项：请参阅大型密集矩阵乘法的循环平铺/阻塞，它在问题中有 5 个嵌套循环，但在答案中有 6 个嵌套循环。 （不过，只是标量，而不是矢量化）。

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

---

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

其他 matmul 策略包括在内循环内添加目标，以便在内循环内加载C和A ，同时提升来自B的负载。 （或者 B 和 A 交换了，我忘记了。） 每个程序员应该知道什么关于内存？ 对于 SSE2 __m128d向量，有一个在附录中以这种方式工作的向量化缓存阻塞示例。 https://www.akkadia.org/drepper/cpumemory.pdf