c++ - 使用 AVX2 有效地计算 std::complex vector 的绝对值

Question

对于某些实时 DSP 应用程序，我需要计算复数值 vector 的绝对值。

简单的实现看起来像这样

computeAbsolute (std::complex<float>* complexSourceVec,
                 float* realValuedDestinationVec,
                 int vecLength)
{
    for (int i = 0; i < vecLength; ++i)
        realValuedDestinationVec[i] = std::abs (complexSourceVec[i]);
}

我想用基于 AVX2 指令的 AVX2 优化版本替换此实现。以这种方式实现它的最有效方法是什么？

注意:源数据是通过我无权访问的 API 交给我的，因此没有机会更改复杂输入 vector 的布局以提高效率。

Answer 1

受到 Dan M 的回答的启发。我首先通过一些调整实现了他的版本:

首先将其更改为使用更宽的 256 位寄存器，然后用 __attribute__((aligned (32))) 能够使用对齐加载

void computeAbsolute1 (const std::complex<float>* cplxIn, float* absOut, const int length)
{
    for (int i = 0; i < length; i += 8)
    {
        float re[8] __attribute__((aligned (32))) = {cplxIn[i].real(), cplxIn[i + 1].real(), cplxIn[i + 2].real(), cplxIn[i + 3].real(), cplxIn[i + 4].real(), cplxIn[i + 5].real(), cplxIn[i + 6].real(), cplxIn[i + 7].real()};
        float im[8] __attribute__((aligned (32))) = {cplxIn[i].imag(), cplxIn[i + 1].imag(), cplxIn[i + 2].imag(), cplxIn[i + 3].imag(), cplxIn[i + 4].imag(), cplxIn[i + 5].imag(), cplxIn[i + 6].imag(), cplxIn[i + 7].imag()};
        __m256 x4 = _mm256_load_ps (re);
        __m256 y4 = _mm256_load_ps (im);
        __m256 b4 = _mm256_sqrt_ps (_mm256_add_ps (_mm256_mul_ps (x4,x4), _mm256_mul_ps (y4,y4)));
        _mm256_storeu_ps (absOut + i, b4);
    }
}

然而，以这种方式手动调整值似乎是一项可以以某种方式加速的任务。现在这是我想出的解决方案，在 clang 编译的快速测试中运行速度提高了 2-3 倍，并进行了全面优化:

#include <complex>
#include <immintrin.h>

void computeAbsolute2 (const std::complex<float>* __restrict cplxIn, float* __restrict absOut, const int length)
{    
    for (int i = 0; i < length; i += 8)
    {
        // load 8 complex values (--> 16 floats overall) into two SIMD registers
        __m256 inLo = _mm256_loadu_ps (reinterpret_cast<const float*> (cplxIn + i    ));
        __m256 inHi = _mm256_loadu_ps (reinterpret_cast<const float*> (cplxIn + i + 4));

        // seperates the real and imaginary part, however values are in a wrong order
        __m256 re = _mm256_shuffle_ps (inLo, inHi, _MM_SHUFFLE (2, 0, 2, 0));
        __m256 im = _mm256_shuffle_ps (inLo, inHi, _MM_SHUFFLE (3, 1, 3, 1));

        // do the heavy work on the unordered vectors
        __m256 abs = _mm256_sqrt_ps (_mm256_add_ps (_mm256_mul_ps (re, re), _mm256_mul_ps (im, im)));

        // reorder values prior to storing
        __m256d ordered = _mm256_permute4x64_pd (_mm256_castps_pd(abs), _MM_SHUFFLE(3,1,2,0));
        _mm256_storeu_ps (absOut + i, _mm256_castpd_ps(ordered));
    }
}

如果没有人提出更快的解决方案，我想我会继续实现

这可以通过 gcc 和 clang ( on the Godbolt compiler explorer ) 高效编译。