c - 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 性能的关键之一。另见 What Every Programmer Should Know About Memory?它在没有转置的附录中有一个缓存阻塞的 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。