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);
    }
  }
}

我似乎不能给出正确的结果。我是不是遗漏了什么？搜索剂量似乎有很大帮助-每个结果要么只做4x4矩阵，magic或一些特殊的魔术，不太容易读和难以理解.

Answer 1

你说得对，你的vB才是问题所在。您正在加载4个连续整数，但mat2[k+0..3][j]并不是连续的。你实际上得到了mat2[k][j+0..3]。

我忘了什么对母校有用。有时，它可以并行地产生4个结果，而不是对每个结果进行水平求和。

转置你的一个输入矩阵工作，并花费O(N^2)。这是值得的，因为它意味着O(N^3) matmul可以使用顺序访问，而您当前的循环结构变得SIMD友好。

还有更好的方法，例如在使用前就将小块转置，所以当您再次读取它们时，它们仍然是L1缓存中的热点。或者在目标行上循环并添加一个结果，而不是为单个或小的行*列点产品积累完整的结果。缓存阻塞，也就是循环平铺，是良好的匹配性能的关键之一。还请参阅What Every Programmer Should Know About Memory?，它在附录中有一个缓存阻塞的SIMD示例，没有转置。

关于优化矩阵乘法，使用SIMD和缓存阻塞已经写了很多.我建议你搜索一下。大多数情况下，如果它可能是在谈论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));
    }
  }
}

一个广播负载(或单独的load+broadcast没有AVX)仍然比一个集合便宜得多。

当前的代码使用4个插入来完成收集，而不是通过对第一个元素使用MOVD来破坏上一个迭代值上的依赖链，这就更糟了。但是，与load + PUNPCKLDQ相比，即使是4个分散元素的最佳集合也是相当糟糕的。更不用说，这使得您的代码需要SSE4.1。

尽管_mm_mullo_epi32仍然需要SSE4.1，而不是扩展的)。

请注意，整数乘吞吐量通常比FP乘差，特别是在Haswell和以后。FMA单元每个32位元素只有24位宽乘子(对于FMA)，所以使用那些用于32x32=>32位整数的单元需要分裂成两个uop。

Answer 2

这个第一个版本是OP发布的，作为对不属于哪里的问题的编辑。把它转移到一个社区-维基的答案只是为了子孙后代。

第一个版本对于性能来说是绝对的垃圾，是矢量化的最糟糕的方式，在最内循环中做一个hsum到标量，并且使用insert_epi32进行手动收集，甚至不是一个4x4转置。

更新： Woho！我终于想出来了。除了我逻辑中的错误(感谢Peter的帮助)，还有_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循环顺序版本。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掩蔽。