使用 SSE 的矩阵向量和矩阵矩阵乘法

Question

我需要编写矩阵-向量和矩阵-矩阵乘法函数，但我无法理解 SSE 命令。

矩阵和向量的维度总是4的倍数。

我设法编写了如下所示的向量-向量乘法函数：

void vector_multiplication_SSE(float* m, float* n, float* result, unsigned const int size)
{
    int i;

    __declspec(align(16))__m128 *p_m = (__m128*)m;
    __declspec(align(16))__m128 *p_n = (__m128*)n;
    __declspec(align(16))__m128 *p_result = (__m128*)result;

    for (i = 0; i < size / 4; ++i)
        p_result[i] = _mm_mul_ps(p_m[i], p_n[i]);

    // print the result
    for (int i = 0; i < size; ++i)
    {
        if (i % 4 == 0) cout << endl;
        cout << result[i] << '\t';
    }
}

现在我正在尝试实现矩阵向量乘法。

这是我目前的情况：

void multiply_matrix_by_vector_SSE(float* m, float* v, float* result, unsigned const int vector_dims)
{
    int i, j;

    __declspec(align(16))__m128 *p_m = (__m128*)m;
    __declspec(align(16))__m128 *p_v = (__m128*)v;
    __declspec(align(16))__m128 *p_result = (__m128*)result;

    for (i = 0; i < vector_dims; i += 4)
    {
        __m128 tmp = _mm_load_ps(&result[i]);
        __m128 p_m_tmp = _mm_load_ps(&m[i]);

        tmp = _mm_add_ps(tmp, _mm_mul_ps(tmp, p_m_tmp));
        _mm_store_ps(&result[i], tmp);

        // another for loop here? 
    }

    // print the result
    for (int i = 0; i < vector_dims; ++i)
    {
        if (i % 4 == 0) cout << endl;
        cout << result[i] << '\t';
    }
}

这个函数看起来完全错误。我的意思是不仅它不能正常工作，而且我似乎也在朝着错误的方向前进。

谁能帮我实现向量-矩阵和矩阵-矩阵乘法？我非常感谢一些示例代码和非常详细的解释

更新

这是我的第 2 次尝试：

它因 Access reading violation 异常而失败，但仍然感觉更接近

void multiply_matrix_by_vector_SSE(float* m, float* v, float* result, unsigned const int vector_dims)
{
    int i, j;

    __declspec(align(16))__m128 *p_m = (__m128*)m;
    __declspec(align(16))__m128 *p_v = (__m128*)v;
    __declspec(align(16))__m128 *p_result = (__m128*)result;

    for (i = 0; i < vector_dims; ++i)
    {
        p_result[i] = _mm_mul_ps(_mm_load_ps(&m[i]), _mm_load_ps1(&v[i]));
    }

    // print the result
    for (int i = 0; i < vector_dims; ++i)
    {
        if (i % 4 == 0) cout << endl;
        cout << result[i] << '\t';
    }
}

更新 2

void multiply_matrix_by_vector_SSE(float* m, float* v, float* result, unsigned const int vector_dims)
{
    int i, j;
    __declspec(align(16))__m128 *p_m = (__m128*)m;
    __declspec(align(16))__m128 *p_v = (__m128*)v;
    __declspec(align(16))__m128 *p_result = (__m128*)result;

    for (i = 0; i < vector_dims; ++i)
    {
        for (j = 0; j < vector_dims * vector_dims / 4; ++j)
        {
            p_result[i] = _mm_mul_ps(p_v[i], p_m[j]);
        }
    }

    for (int i = 0; i < vector_dims; ++i)
    {
        if (i % 4 == 0) cout << endl;
        cout << result[i] << '\t';
    }
    cout << endl;
}

Answer 1

没有任何技巧或任何东西，矩阵向量乘法只是向量和矩阵的一行之间的一堆点积。您的代码实际上没有那种结构。实际上将其写为点积（未测试）：

for (int row = 0; row < nrows; ++row) {
    __m128 acc = _mm_setzero_ps();
    // I'm just going to assume the number of columns is a multiple of 4
    for (int col = 0; col < ncols; col += 4) {
        __m128 vec = _mm_load_ps(&v[col]);
        // don't forget it's a matrix, do 2d addressing
        __m128 mat = _mm_load_ps(&m[col + ncols * row]);
        acc = _mm_add_ps(acc, _mm_mul_ps(mat, vec));
    }
    // now we have 4 floats in acc and they have to be summed
    // can use two horizontal adds for this, they kind of suck but this
    // isn't the inner loop anyway.
    acc = _mm_hadd_ps(acc, acc);
    acc = _mm_hadd_ps(acc, acc);
    // store result, which is a single float
    _mm_store_ss(&result[row], acc);
}

有一些明显的技巧，例如一次处理多行、重用向量的负载以及创建多个独立的依赖链以便更好地利用吞吐量（见下文）。还有一个非常简单的技巧是将 FMA 用于 mul/add 组合，但支持还没有那么广泛（2015 年还没有，但现在在 2020 年相当普遍）。

您可以由此构建矩阵-矩阵乘法（如果您更改结果所在的位置），但这不是最优的（请参阅下文）。

一次取四行（未测试）：

for (int row = 0; row < nrows; row += 4) {
    __m128 acc0 = _mm_setzero_ps();
    __m128 acc1 = _mm_setzero_ps();
    __m128 acc2 = _mm_setzero_ps();
    __m128 acc3 = _mm_setzero_ps();
    for (int col = 0; col < ncols; col += 4) {
        __m128 vec = _mm_load_ps(&v[col]);
        __m128 mat0 = _mm_load_ps(&m[col + ncols * row]);
        __m128 mat1 = _mm_load_ps(&m[col + ncols * (row + 1)]);
        __m128 mat2 = _mm_load_ps(&m[col + ncols * (row + 2)]);
        __m128 mat3 = _mm_load_ps(&m[col + ncols * (row + 3)]);
        acc0 = _mm_add_ps(acc0, _mm_mul_ps(mat0, vec));
        acc1 = _mm_add_ps(acc1, _mm_mul_ps(mat1, vec));
        acc2 = _mm_add_ps(acc2, _mm_mul_ps(mat2, vec));
        acc3 = _mm_add_ps(acc3, _mm_mul_ps(mat3, vec));
    }
    acc0 = _mm_hadd_ps(acc0, acc1);
    acc2 = _mm_hadd_ps(acc2, acc3);
    acc0 = _mm_hadd_ps(acc0, acc2);
    _mm_store_ps(&result[row], acc0);
}

现在每 4 个 FMA 只有 5 个负载，而在未展开行的版本中，每 1 个 FMA 有 2 个负载。还有 4 个独立的 FMA，或 add/mul 对，没有 FMA 收缩，无论哪种方式，它都会增加 pipelined/simultaneous 执行的可能性。实际上您可能想要展开更多，例如 Skylake 每个周期可以启动 2 个独立的 FMA，它们需要 4 个周期才能完成，因此要完全占用两个 FMA 单元，您需要 8 个独立的 FMA。作为奖励，对于水平求和，这 3 个水平加法最终的结果相对较好。

不同的数据布局最初似乎是一个缺点，不再可能简单地从矩阵和向量中进行向量加载并将它们相乘（这会将第一个矩阵的一个小行向量乘以tiny row 第二个矩阵的向量，这是错误的）。但是全矩阵-矩阵乘法可以利用这样一个事实，即它本质上是将一个矩阵乘以许多独立的向量，它充满了独立的工作要做。水平和也可以很容易地避免。所以实际上它比矩阵向量乘法更方便。

关键是从矩阵 A 中取出一个小列向量和从矩阵 B 中取出一个小行向量，并将它们相乘成一个小矩阵。与您习惯的做法相比，这听起来可能是相反的，但这样做对于 SIMD 效果更好，因为计算始终保持独立且无水平操作。

例如（未测试，假设矩阵的维度可以被展开因子整除，需要 x64 否则它会用完寄存器）

for (size_t i = 0; i < mat1rows; i += 4) {
    for (size_t j = 0; j < mat2cols; j += 8) {
        float* mat1ptr = &mat1[i * mat1cols];
        __m256 sumA_1, sumB_1, sumA_2, sumB_2, sumA_3, sumB_3, sumA_4, sumB_4;
        sumA_1 = _mm_setzero_ps();
        sumB_1 = _mm_setzero_ps();
        sumA_2 = _mm_setzero_ps();
        sumB_2 = _mm_setzero_ps();
        sumA_3 = _mm_setzero_ps();
        sumB_3 = _mm_setzero_ps();
        sumA_4 = _mm_setzero_ps();
        sumB_4 = _mm_setzero_ps();

        for (size_t k = 0; k < mat2rows; ++k) {
            auto bc_mat1_1 = _mm_set1_ps(mat1ptr[0]);
            auto vecA_mat2 = _mm_load_ps(mat2 + m2idx);
            auto vecB_mat2 = _mm_load_ps(mat2 + m2idx + 4);
            sumA_1 = _mm_add_ps(_mm_mul_ps(bc_mat1_1, vecA_mat2), sumA_1);
            sumB_1 = _mm_add_ps(_mm_mul_ps(bc_mat1_1, vecB_mat2), sumB_1);
            auto bc_mat1_2 = _mm_set1_ps(mat1ptr[N]);
            sumA_2 = _mm_add_ps(_mm_mul_ps(bc_mat1_2, vecA_mat2), sumA_2);
            sumB_2 = _mm_add_ps(_mm_mul_ps(bc_mat1_2, vecB_mat2), sumB_2);
            auto bc_mat1_3 = _mm_set1_ps(mat1ptr[N * 2]);
            sumA_3 = _mm_add_ps(_mm_mul_ps(bc_mat1_3, vecA_mat2), sumA_3);
            sumB_3 = _mm_add_ps(_mm_mul_ps(bc_mat1_3, vecB_mat2), sumB_3);
            auto bc_mat1_4 = _mm_set1_ps(mat1ptr[N * 3]);
            sumA_4 = _mm_add_ps(_mm_mul_ps(bc_mat1_4, vecA_mat2), sumA_4);
            sumB_4 = _mm_add_ps(_mm_mul_ps(bc_mat1_4, vecB_mat2), sumB_4);
            m2idx += 8;
            mat1ptr++;
        }
        _mm_store_ps(&result[i * mat2cols + j], sumA_1);
        _mm_store_ps(&result[i * mat2cols + j + 4], sumB_1);
        _mm_store_ps(&result[(i + 1) * mat2cols + j], sumA_2);
        _mm_store_ps(&result[(i + 1) * mat2cols + j + 4], sumB_2);
        _mm_store_ps(&result[(i + 2) * mat2cols + j], sumA_3);
        _mm_store_ps(&result[(i + 2) * mat2cols + j + 4], sumB_3);
        _mm_store_ps(&result[(i + 3) * mat2cols + j], sumA_4);
        _mm_store_ps(&result[(i + 3) * mat2cols + j + 4], sumB_4);
    }
}

该代码的要点是很容易安排计算非常 SIMD 友好，有很多独立的算法来饱和浮点单元，同时使用相对较少的负载（这否则可能成为瓶颈，即使不考虑它们可能会错过 L1 缓存，只是因为它们太多了）。

您甚至可以使用此代码，但它无法与 Intel MKL 竞争。特别是对于中型或大型矩阵，平铺非常重要。将它升级到 AVX 很容易。它根本不适合小矩阵，例如乘以两个 4x4 矩阵见 Efficient 4x4 matrix multiplication。

使用 SSE 的矩阵向量和矩阵矩阵乘法

Matrix-Vector and Matrix-Matrix multiplication using SSE

c++

sse

intrinsics

matrix-multiplication

vector-multiplication

更新

更新 2