SSE矩阵 - 矩阵乘法

时间:2016-10-28 21:29:12

标签: c sse matrix-multiplication

我在使用SS中的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矩阵,mat-vec或一些不易阅读且难以理解的特殊魔法......

更新 Woho!我终于弄明白了。除了我的逻辑中的错误(感谢Peter Cordes的帮助),还有_mm_mul_epi32()的问题没有像我想的那样工作 - 我应该使用_mm_mullo_epi32()代替!

我知道这不是最有效的代码,但它是为了让它首先正常工作 - 现在我可以继续优化它。

void matmulSSE(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]);
                vB = _mm_insert_epi32(vB, mat2[k][j], 0);
                vB = _mm_insert_epi32(vB, mat2[k + 1][j], 1);
                vB = _mm_insert_epi32(vB, mat2[k + 2][j], 2);
                vB = _mm_insert_epi32(vB, mat2[k + 3][j], 3);
                vR = _mm_mullo_epi32(vA, vB);
                vR = _mm_hadd_epi32(vR, vR);
                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");
            }
        }
    }
}

更新2:将Peters示例转换为i-k-j循环订单版本。需要为vR额外加载并在存储中移动到内部循环,但是设置vA可以向上移动一个循环。结果更快。

void matmulSSE_2(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");
            }
        }
    }
}

2 个答案:

答案 0 :(得分: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,但它也适用于整数。

(除了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来破坏前一个迭代值的依赖关系链,因此更糟糕。但是,即使是4个分散元素的最佳聚集与load + PSHUFD相比也是相当糟糕的。更不用说那使你的代码需要SSE4.1。虽然它仍然适用于_mm_mullo_epi32,而不是扩展PMULDQ (_mm_mul_epi32)

答案 1 :(得分:0)

这是由OP发布的,以编辑不属于该问题的问题。


更新: 哇!我终于弄明白了。除了我的逻辑中的错误(感谢彼得·科德斯的帮助)之外,还有_mm_mul_epi32()不能按我认为的那样工作的问题-我应该使用_mm_mullo_epi32()来代替!

我知道这不是最有效的代码,但它是为了使它首先正常工作而设计的-现在,我可以继续对其进行优化。

// editor's note: this is the most naive way to vectorize
void matmulSSE(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]);
                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);
                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");
            }
        }
    }
}

更新2:将Peters示例转换为i-k-j循环订单版本。 vR需要额外的负载,并且需要在存储中移动到内部循环,但是可以将设置vA向上移动一个循环。结果更快。

// this is significantly better but doesn't do any cache-blocking
void matmulSSE_2(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");
            }
        }
    }
}