SSE矩陣-矩陣乘法

Question

我在用 C 中的 SSE 進行矩陣-矩陣乘法時遇到問題。

這是我到目前為止所得到的：

#define N 1000

void matmulSSE(int mat1[N][N], int mat2[N][N], int result[N][N]) {
  int i, j, k;
  __m128i vA, vB, vR;

  for(i = 0; i < N; ++i) {
    for(j = 0; j < N; ++j) {
        vR = _mm_setzero_si128();
        for(k = 0; k < N; k += 4) {
            //result[i][j] += mat1[i][k] * mat2[k][j];
            vA = _mm_loadu_si128((__m128i*)&mat1[i][k]);
            vB = _mm_loadu_si128((__m128i*)&mat2[k][j]); //how well does the k += 4 work here? Should it be unrolled?
            vR = _mm_add_epi32(vR, _mm_mul_epi32(vA, vB));
        }
        vR = _mm_hadd_epi32(vR, vR);
        vR = _mm_hadd_epi32(vR, vR);
        result[i][j] += _mm_extract_epi32(vR, 0);
    }
  }
}

我似乎無法讓它給出正確的結果。 我錯過了什么嗎？ 搜索 dosent 似乎有很大幫助 - 每個結果要么只做 4x4 矩陣，mat-vec 或一些不太可讀且難以理解的特殊魔法......

Answer 1

你是對的，你的vB是問題所在。 您正在加載 4 個連續的整數，但mat2[k+0..3][j]不連續。 你實際上得到了mat2[k][j+0..3] 。

---

我忘了什么對 matmul 有效。 有時並行產生 4 個結果效果很好，而不是對每個結果進行水平求和。

轉置您的輸入矩陣之一有效，並且成本為 O(N^2)。 這是值得的，因為這意味着 O(N^3) matmul 可以使用順序訪問，並且您當前的循環結構變得對 SIMD 友好。

還有更好的方法，比如在使用前轉置小塊，這樣當你再次讀取它們時，它們在 L1 緩存中仍然很熱。 或者循環遍歷目標行並添加一個結果，而不是為單個或一小組行*列點積累積完整結果。 緩存阻塞，又名循環平鋪，是良好 matmul 性能的關鍵之一。 另請參閱每個程序員應該了解的關於內存的內容？ 它在沒有轉置的附錄中有一個緩存阻塞的 SIMD FP matmul 示例。

關於優化矩陣乘法、使用 SIMD 和緩存阻塞的文章很多。 建議你google一下。 大多數可能是在談論 FP，但它也適用於整數。

（除了 SSE/AVX 只有 FP 的 FMA，沒有 32 位整數，並且 8 位和 16 位輸入 PMADD 指令執行對的水平加法。）

---

實際上我認為你可以在這里並行產生 4 個結果，如果一個輸入已經被轉置了：

void matmulSSE(int mat1[N][N], int mat2[N][N], int result[N][N]) {

  for(int i = 0; i < N; ++i) {
    for(int j = 0; j < N; j+=4) {   // vectorize over this loop
        __m128i vR = _mm_setzero_si128();
        for(int k = 0; k < N; k++) {   // not this loop
            //result[i][j] += mat1[i][k] * mat2[k][j];
            __m128i vA = _mm_set1_epi32(mat1[i][k]);  // load+broadcast is much cheaper than MOVD + 3 inserts (or especially 4x insert, which your new code is doing)
            __m128i vB = _mm_loadu_si128((__m128i*)&mat2[k][j]);  // mat2[k][j+0..3]
            vR = _mm_add_epi32(vR, _mm_mullo_epi32(vA, vB));
        }
        _mm_storeu_si128((__m128i*)&result[i][j], vR));
    }
  }
}

廣播加載（或單獨的加載+廣播，沒有 AVX）仍然比收集便宜得多。

您當前的代碼使用 4 個插入進行收集，而不是通過對第一個元素使用 MOVD 來破壞對上一次迭代值的依賴鏈，因此情況更糟。 但與負載 + PUNPCKLDQ 相比，即使是最好的 4 個分散元素的集合也很糟糕。 更不用說這使您的代碼需要 SSE4.1。

盡管_mm_mullo_epi32無論如何它都需要 SSE4.1 而不是加寬的PMULDQ ( _mm_mul_epi32 ) 。

請注意，整數乘法吞吐量通常比 FP 乘法差，尤其是在 Haswell 及更高版本上。 FP FMA 單元每個 32 位元素只有 24 位寬的乘法器（對於 FP 尾數），因此使用那些用於 32x32=>32 位整數的乘法器需要分成兩個 uops。

Answer 2

如果轉置第二個矩陣，則第一種方法是正確的

Answer 3

第一個版本由 OP 發布，作為對不屬於它的問題的編輯。 將其移至社區維基答案，僅供后人使用。

第一個版本對於性能來說絕對是垃圾，最糟糕的矢量化方式，在最內部循環中將 hsum 降為標量，並使用insert_epi32進行手動收集，甚至不是 4x4 轉置。

---

更新：哇！ 我終於想通了。 除了我的邏輯錯誤（感謝 Peter Cordes 的幫助）之外，還有_mm_mul_epi32()沒有像我想象的那樣工作 - 我應該使用_mm_mullo_epi32()來代替！

我知道這不是最有效的代碼，但它是為了讓它首先正常工作 - 現在我可以繼續優化它。

（注意，不要使用這個，它非常非常慢）

// editor's note: this is the most naive and horrible way to vectorize
void matmulSSE_inefficient(int mat1[N][N], int mat2[N][N], int result[N][N]) {
    int i, j, k;
    __m128i vA, vB, vR, vSum;

    for(i = 0; i < N; ++i) {
        for(j = 0; j < N; ++j) {
            vR = _mm_setzero_si128();
            for(k = 0; k < N; k += 4) {
                //result[i][j] += mat1[i][k] * mat2[k][j];
                vA = _mm_loadu_si128((__m128i*)&mat1[i][k]);
                          // less braindead would be to start vB with movd, avoiding a false dep and one shuffle uop.
                          // vB = _mm_cvtsi32_si128(mat2[k][j]);   // but this manual gather is still very bad
                vB = _mm_insert_epi32(vB, mat2[k][j], 0);     // false dependency on old vB
                vB = _mm_insert_epi32(vB, mat2[k + 1][j], 1);  // bad spatial locality
                vB = _mm_insert_epi32(vB, mat2[k + 2][j], 2);  // striding down a column
                vB = _mm_insert_epi32(vB, mat2[k + 3][j], 3);
                vR = _mm_mullo_epi32(vA, vB);
                vR = _mm_hadd_epi32(vR, vR);                // very slow inside the inner loop
                vR = _mm_hadd_epi32(vR, vR);
                result[i][j] += _mm_extract_epi32(vR, 0);

                //DEBUG
                //printf("vA: %d, %d, %d, %d\n", vA.m128i_i32[0], vA.m128i_i32[1], vA.m128i_i32[2], vA.m128i_i32[3]);
                //printf("vB: %d, %d, %d, %d\n", vB.m128i_i32[0], vB.m128i_i32[1], vB.m128i_i32[2], vB.m128i_i32[3]);
                //printf("vR: %d, %d, %d, %d\n", vR.m128i_i32[0], vR.m128i_i32[1], vR.m128i_i32[2], vR.m128i_i32[3]);
                //printf("\n");
            }
        }
    }
}

最初由 OP 編寫的極其低效的代碼結束

---

更新 2：將 OP 的示例轉換為 ikj 循環順序版本。 需要額外的 vR 負載並將存儲移動到內循環中，但設置 vA 可以向上移動一個循環。 結果更快。

// this is significantly better but doesn't do any cache-blocking
void matmulSSE_v2(int mat1[N][N], int mat2[N][N], int result[N][N]) {
    int i, j, k;
    __m128i vA, vB, vR;

    for(i = 0; i < N; ++i) {
        for(k = 0; k < N; ++k) {
            vA = _mm_set1_epi32(mat1[i][k]);
            for(j = 0; j < N; j += 4) {
                //result[i][j] += mat1[i][k] * mat2[k][j];
                vB = _mm_loadu_si128((__m128i*)&mat2[k][j]);
                vR = _mm_loadu_si128((__m128i*)&result[i][j]);
                vR = _mm_add_epi32(vR, _mm_mullo_epi32(vA, vB));
                _mm_storeu_si128((__m128i*)&result[i][j], vR);

                //DEBUG
                //printf("vA: %d, %d, %d, %d\n", vA.m128i_i32[0], vA.m128i_i32[1], vA.m128i_i32[2], vA.m128i_i32[3]);
                //printf("vB: %d, %d, %d, %d\n", vB.m128i_i32[0], vB.m128i_i32[1], vB.m128i_i32[2], vB.m128i_i32[3]);
                //printf("vR: %d, %d, %d, %d\n", vR.m128i_i32[0], vR.m128i_i32[1], vR.m128i_i32[2], vR.m128i_i32[3]);

                //printf("\n");
            }
        }
    }
}

---

這些假設 N 是 4 的倍數，向量寬度

如果不是這種情況，通常將數組存儲填充到向量寬度的倍數通常會更容易，因此每行末尾都有填充，您可以使用簡單的j < N; j += 4 j < N; j += 4循環條件。 您需要使用 4 或 8 的倍數的行步長來跟蹤與存儲布局分開的實際N大小。

否則你需要一個循環條件，比如j < N-3 ; j += 4`，以及行尾的標量清理。

或者將最后一個完整向量屏蔽或保留在寄存器中，以便您可以_mm_alignr_epi8使用可能重疊的最終向量，該向量在行的末尾結束，並且可以進行向量存儲。 使用 AVX 或特別是 AVX512 掩蔽更容易。