使用 AVX2 高效计算 std::complex<float> 向量的绝对值

Question

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

简单的实现看起来像这样

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 交给我的，因此没有机会更改复杂输入向量的布局以提高效率。

Answer 1

很难（如果可能的话）编写复数 abs 的 "highly optimized AVX2" 版本，因为标准中定义复数的方式可以防止（特别是由于所有 inf/nan 极端情况）很多优化。

然而，如果您不关心正确性，您可以使用 -ffast-math，一些编译器会为您优化代码。查看 gcc 输出：https://godbolt.org/z/QbZlBI

您也可以使用此输出并使用内联汇编创建您自己的 abs 函数。但是，是的，正如已经提到的，如果您真的需要性能，您可能想要将 std::complex 换成其他东西。

通过手动填充小型 re 和 im 数组，我能够为您的特定案例获得不错的输出，并进行所有必需的随机播放。参见：https://godbolt.org/z/sWAAXo 这可以简单地扩展为 ymm 个寄存器。

无论如何，这是改编自的最终解决方案，它使用内部函数并结合巧妙的编译器优化：

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

static inline void cabs_soa4(const float *re, const float *im, float *b) {
    __m128 x4 = _mm_loadu_ps(re);
    __m128 y4 = _mm_loadu_ps(im);
    __m128 b4 = _mm_sqrt_ps(_mm_add_ps(_mm_mul_ps(x4,x4), _mm_mul_ps(y4,y4)));
    _mm_storeu_ps(b, b4);
}

void computeAbsolute (const std::complex<float>* src,
                 float* realValuedDestinationVec,
                 int vecLength)
{
    for (int i = 0; i < vecLength; i += 4) {
        float re[4] = {src[i].real(), src[i + 1].real(), src[i + 2].real(), src[i + 3].real()};
        float im[4] = {src[i].imag(), src[i + 1].imag(), src[i + 2].imag(), src[i + 3].imag()};
        cabs_soa4(re, im, realValuedDestinationVec);
    }
}

编译为简单

_Z15computeAbsolutePKSt7complexIfEPfi:
        test    edx, edx
        jle     .L5
        lea     eax, [rdx-1]
        shr     eax, 2
        sal     rax, 5
        lea     rax, [rdi+32+rax]
.L3:
        vmovups xmm0, XMMWORD PTR [rdi]
        vmovups xmm2, XMMWORD PTR [rdi+16]
        add     rdi, 32
        vshufps xmm1, xmm0, xmm2, 136
        vmulps  xmm1, xmm1, xmm1
        vshufps xmm0, xmm0, xmm2, 221
        vfmadd132ps     xmm0, xmm1, xmm0
        vsqrtps xmm0, xmm0
        vmovups XMMWORD PTR [rsi], xmm0
        cmp     rax, rdi
        jne     .L3
.L5:
        ret

https://godbolt.org/z/Yu64Wg

Answer 2

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

首先将其更改为使用更宽的 256 位寄存器，然后将临时 re 和 im 数组标记为 __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) 高效编译。

使用 AVX2 高效计算 std::complex<float> 向量的绝对值

Efficiently compute absolute values of std::complex<float> vector with AVX2

c++

avx2