使用 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