将SSE矩阵矢量乘法代码转换为AVX

Convert SSE matrix-vector multiplication code to AVX

本文关键字：代码 转换 AVX SSE      更新时间：2023-10-16

我正在尝试将SSE函数转换为AVX。这个函数做向量矩阵乘法，这是我的SSE代码：

void multiply_matrix_by_vector_SSE(float* m, float* v, float* result, unsigned const int vector_dims)
{
    size_t i, j;
    for (i = 0; i < vector_dims; ++i)
    {
        __m128 acc = _mm_setzero_ps();
        for (j = 0; j < vector_dims; j += 4)
        {
            __m128 vec = _mm_load_ps(&v[j]);
            __m128 mat = _mm_load_ps(&m[j + vector_dims * i]);
            //acc = _mm_add_ps(acc, _mm_mul_ps(mat, vec));
            acc = _mm_fmadd_ps(mat, vec, acc);
        }
        acc = _mm_hadd_ps(acc, acc);
        acc = _mm_hadd_ps(acc, acc);
        _mm_store_ss(&result[i], acc);
    }
}

以下是我对AVX的看法：

void multiply_matrix_by_vector_AVX(float* m, float* v, float* result, unsigned const int vector_dims)
{
    size_t i, j;
    for (i = 0; i < vector_dims; ++i)
    {
        __m256 acc = _mm256_setzero_ps();
        for (j = 0; j < vector_dims; j += 8)
        {
            __m256 vec = _mm256_load_ps(&v[j]);
            __m256 mat = _mm256_load_ps(&m[j + vector_dims * i]);
            acc = _mm256_fmadd_ps(mat, vec, acc);
        }
        acc = _mm256_hadd_ps(acc, acc);
        acc = _mm256_hadd_ps(acc, acc);
        acc = _mm256_hadd_ps(acc, acc);
        acc = _mm256_hadd_ps(acc, acc);
        _mm256_store_ps(&result[i], acc);
    }
}

但是AVX代码崩溃（Access violation reading location 0xFFFFFFFFFFFFFFFF）。

有人能帮我使我的AVX功能正常工作吗？

PS：我在函数中传递的矩阵和向量的大小总是8的倍数。此外，我传递给SSE函数的数组是16位对齐的（__declspec(align(16))float* = generate_matrix(256);），而我传递给AVX函数的数组则是32位对齐（__declspec(align(32))float* = generate_matrix(256);）；

不幸的是，使用这样的水平加法并不能简单地扩展到256位，因为指令（以及大多数其他指令）是"laned"的——它的作用就像两个并行的haddps，一个在上半部分，一个位于下半部分，没有混合，所以下半部分和上半部分不会相加。

当然，这仍然不是一个压缩的结果，并且该压缩的存储中有一个对齐的存储写入某个未对齐的地址，并且将失败（这个错误有点奇怪，但不管怎样）。

无论如何，让我们修复水平总和：（未测试）

// this part still works
acc = _mm256_hadd_ps(acc, acc);
acc = _mm256_hadd_ps(acc, acc);
// this is new
__m128 acc1 = _mm256_extractf128_ps(acc, 0);
__m128 acc2 = _mm256_extractf128_ps(acc, 1);
acc1 = _mm_add_ss(acc1, acc2);
// do scalar store, obviously
_mm_store_ss(&result[i], acc1);

顺便说一句，内环需要10个独立的链（和10个累加器）才能最大限度地提高Haswell的吞吐量。

相关文章：