c++ - 使用 SSE 计算矩阵乘积比使用直接算法慢得多

Question

我想使用直接算法将两个矩阵相乘一次:

template <typename T>
void multiplicate_straight(T ** A, T ** B, T ** C, int sizeX)
{
    T ** D = AllocateDynamicArray2D<T>(sizeX, sizeX);
    transpose_matrix(B, D,sizeX);
    for(int i = 0; i < sizeX; i++)
    {
        for(int j = 0; j < sizeX; j++)
        {
            for(int g = 0; g < sizeX; g++)
            {
                C[i][j] += A[i][g]*D[j][g];
            }
        }
    }
    FreeDynamicArray2D<T>(D);
}

还有一次是通过使用 SSE 函数。为此，我创建了两个函数:

template <typename T>
void SSE_vectormult(T * A, T * B, int size)
{

    __m128d a;
    __m128d b;
    __m128d c;
#ifdef linux
    double A2[2], B2[2], C[2] __attribute__ ((aligned(16)));
#endif
#ifdef _WIN32
    __declspec(align(16)) double A2[2], B2[2], C[2];
#endif
    for(int i = 0; i < size; i+=2)
    {
        //std::cout << "In SSE_vectormult: i is: " << i << '\n';
        A2[0] = A[i];
        B2[0] = B[i];
        A2[1] = A[i+1];
        B2[1] = B[i+1];
        //std::cout << "Values from A and B written to A2 and B2\n";
        a = _mm_load_pd(A2);
        b = _mm_load_pd(B2);
        //std::cout << "Values converted to a and b\n";
        c = _mm_mul_pd(a,b);
        _mm_store_pd(C, c);
        A[i] = C[0];
        A[i+1] = C[1];
    };
}

和

template <typename T>
void multiplicate_SSE(T ** A, T ** B, T ** C, int sizeX)
{
//    std::cout << "Entered SSE-Function\n";
    T ** D = AllocateDynamicArray2D<T>(sizeX, sizeX);
    T * tmp = AllocateDynamicArray1D<T>(sizeX);
    T * tmp2 = AllocateDynamicArray1D<T>(sizeX);
    //std::cout << "Matrices allocated\n";
    transpose_matrix<T>(B, D,sizeX);
    //std::cout << "Matrix B transposed\n";
    for(int i = 0; i < sizeX; i++)
    {
        for(int j = 0; j < sizeX; j++)
        {
            extract_row<T>(A,tmp, i, sizeX);
//            std::cout << "Row from A extracted\n";
            //print_vector(tmp, sizeX);
            extract_row<T>(D, tmp2, j, sizeX);
//            std::cout << "Row from D extracted\n";
            //print_vector(tmp2, sizeX);
            SSE_vectormult<T>(tmp, tmp2, sizeX);
//            std::cout << "Vectors multiplicated\n";
            //print_vector(tmp, sizeX);
            C[i][j] = add_vector(tmp, sizeX);
//            std::cout << "Written value to C\n";
//            std::cout << "j is " << j << " and i is " << i << '\n';
        }
    }
//    std::cout << "Loop finished\n";
    FreeDynamicArray2D<T>(D);
    //std::cout << "Freed D\n";
    //FreeDynamicArray1D<T>(tmp);????
//    std::cout << "Freed tmp\n";
    FreeDynamicArray1D<T>(tmp2);
//    std::cout << "Everything freed, returning\n";
}

但随后我遇到了几个问题:一方面，当我想释放 multiplicate_SSE() 中标有几个问号的 tmp 数组时，我收到错误“_BLOCK_TYPE_IS_VALID”。我考虑过两次释放相同空间的可能性，所以我取消了注释(但我想我会因此发生内存泄漏吗？)。现在，当我用相同的矩阵比较这两个函数的性能时，对于两个 1024x1024 矩阵，SSE 函数需要比直接方法长四倍左右的时间。
如何重写 SSE 函数以获得更好的性能(我以前从未使用过 SSE)，以及如何修复内存泄漏？
谢谢!

Answer 1

通过在标量代码中进行转置，您的想法是正确的，但在使用 SSE 时您并不需要完全转置。

让我们坚持 float (SGEMM)。你想用 SSE 做的是一次做四个点积。你想要 C = A*B。我们来看一个 8x8 矩阵。假设 B 是:

(0   1  2  3) ( 4  5  6  7)
(8   9 10 11) (12 13 14 15) 
(16 17 18 19) (20 21 22 23)
(24 25 26 27) (28 29 30 31)
(32 33 34 35) (36 37 38 39)
(40 41 42 43) (44 45 46 47)
(48 49 50 51) (52 53 54 55)
(56 57 58 59) (60 61 62 63)

所以对于 SSE，您可以:

C[0][0] C[0][1] C[0][2] C[0][3] = 
A[0][0]*(0 1 2 3) + A[0][1]*(8 9 10 11) + A[0][2]*(16 17 18 19)...+ A[0][7]*(56 57 58 59)

这一次得到四个点积。问题是您必须在 B 中向下移动一列，并且这些值不在同一缓存行中。 如果宽度为四的每一列在内存中是连续的就更好了。因此，不是对每个元素进行转置，而是像这样转置宽度为 4 的 strip :

(0  1  2  3)( 8  9 10 11)(16 17 18 19)(24 25 26 27)(32 33 34 35)(40 41 42 43)(48 49 50 51)(56 57 58 59)
(4  5  6  7)(12 13 14 15)(20 21 22 23)(28 29 30 31)(36 37 38 39)(44 45 46 47)(52 53 54 55)(60 61 62 63)

如果您将括号中的四个值中的每一个都视为一个单元，则这相当于将 8x2 矩阵转置为 2x8 矩阵。现在请注意 B 的宽度为 4 的列 在内存中是连续的。这对缓存更友好。对于 8x8 矩阵，这不是真正的问题，但例如对于 1024x1024 矩阵，它就是问题。请参阅下面的代码了解如何执行此操作。对于 AVX，您转置宽度为 8 的 strip (这意味着您无需为 8x8 矩阵做任何事情)。对于双宽度，SSE 为 2，AVX 为 4。

假设矩阵适合缓存，这应该比标量代码快四倍。但是，对于大型矩阵，此方法仍将受内存限制，因此您的 SSE 代码可能不会比标量代码快很多(但应该不会更差)。

但是，如果您使用循环平铺和重新排列矩阵(适合 L2 缓存)而不是整个矩阵矩阵乘法是计算限制而不是内存限制，即使对于不需要的非常大的矩阵也是如此不适合 L3 缓存。 那是另一个话题。

编辑:一些(未经测试的)代码与您的标量代码进行比较。我将循环展开 2。

void SGEMM_SSE(const float *A, const float *B, float *C, const int sizeX) {
    const int simd_width = 4;
    const int unroll = 2;
    const int strip_width = simd_width*unroll
    float *D = (float*)_mm_malloc(sizeof(float)*sizeX*sizeX, 16);
    transpose_matrix_strip(B, D,sizeX, strip_width); //tranpose B in strips of width eight
    for(int i = 0; i < sizeX; i++) {
        for(int j = 0; j < sizeX; j+=strip_width) {
            float4 out_v1 = 0; //broadcast (0,0,0,0)
            float4 out_V2 = 0;
            //now calculate eight dot products
            for(int g = 0; g < sizeX; g++) {
                //load eight values rrom D into two SSE registers
                float4 vec4_1.load(&D[j*sizeX + strip_width*g]);
                float4 vec4_2.load(&D[j*sizeX + strip_width*g + simd_width]);
                out_v1 += A[i][g]*vec4_v1;
                out_v2 += A[i][g]*vec4_v2;
            }
            //store eight dot prodcuts into C
            out_v1.store(&C[i*sizeX + j]);
            out_v2.store(&C[i*sizeX + j + simd_width]);
        }
    }
    _mm_free(D);
}

void transpose_matrix_strip(const float* A, float* B, const int N, const int strip_width) {
    //#pragma omp parallel for
    for(int n=0; n<N*N; n++) {
        int k = strip_width*(n/N/strip_width);
        int i = (n/strip_width)%N;
        int j = n%strip_width;
        B[n] = A[N*i + k + j];
    }
}

注意 j 现在递增 8。更多展开可能会有所帮助。如果您想使用内部函数，您可以使用 _mm_load_ps、_mm_store_ps、_mm_set1_ps(对于广播，例如 _mm_set1_ps(A[i][ g]))、_mm_add_ps 和 _mm_mul_ps。就是这样。