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 友好。

还有更好的方法，比如在使用前转置小块，这样当你再次读取它们时它们在一级缓存中仍然很热。或者遍历目标行并添加一个结果，而不是为单个或一小组行*列点积累积完整结果。缓存阻塞，又名循环平铺，是良好 matmul 性能的关键之一。另见 What Every Programmer Should Know About Memory?，附录中有一个缓存阻塞的 SIMD FP matmul 示例，没有转置。

关于使用 SIMD 和缓存阻塞优化矩阵乘法的文章很多。我建议你 google 起来。大多数如果它可能在谈论 FP，但它也适用于整数。

（除了 SSE/AVX 仅对 FP 有 FMA，对 32 位整数没有 FMA，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。

尽管它无论如何都需要 SSE4.1 _mm_mullo_epi32 而不是扩展 PMULDQ (_mm_mul_epi32).

请注意，整数乘法吞吐量通常比 FP 乘法差，尤其是在 Haswell 和更高版本上。 FP FMA 单元每个 32 位元素（对于 FP 尾数）只有 24 位宽的乘法器，因此将这些乘法器用于 32x32=>32 位整数需要分成两个 uops。

Answer 2

第一个版本由 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 示例转换为 i-k-j 循环顺序版本。需要为 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 循环健康）状况。您需要跟踪实际 N 大小与存储布局分开，行步幅是 4 或 8 的倍数。

否则你想要一个像j < N-3这样的循环条件； j += 4`，以及一行末尾的标量清理。

或者将最后一个完整向量屏蔽或保留在寄存器中，这样您就可以 _mm_alignr_epi8 使用在行末尾结束的可能重叠的最终向量，并可能进行向量存储。使用 AVX 或尤其是 AVX512 掩码会更容易。

SSE 矩阵-矩阵乘法

SSE matrix-matrix multiplication

c

sse

matrix-multiplication