Efficiently compute absolute values of std::complex vector with AVX2

Question

For some real-time DSP application I need to compute the absolute values of a complex valued vector.

The straightforward implementation would look like that

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

I want to replace this implementation with an AVX2 optimized version, based on AVX2 instrincts. What would be the most efficient way to implement it that way?

Note: The source data is handed to me by an API I have no access to, so there is no chance to change the layout of the complex input vector for better efficiency.

Answer 1

Inspired by the answer of Dan M. I first implemented his version with some tweaks:

First changed it to use the wider 256 Bit registers, then marked the temporary re and im arrays with __attribute__((aligned (32))) to be able to use aligned load

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);
    }
}

However manually shuffling the values this way seemed like a task that could be speeded up somehow. Now this is the solution I came up with, that runs 2 - 3 times faster in a quick test compiled by clang with full optimization:

#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));
    }
}

I think I'll go with that implementation if no one comes up with a faster solution

This compiles efficiently with gcc and clang (on the Godbolt compiler explorer).

Answer 2

It's really hard (if possible) to write "highly optimized AVX2" version of complex abs since the way complex numbers are defined in the standard prevents (specifically due to all inf/nan corner cases) a lot of optimization.

However, if you don't care about the correctness you can just use -ffast-math and some compilers would optimize the code for you. See gcc output: https://godbolt.org/z/QbZlBI

You can also take this output and create your own abs function with inline assembly. But yes, as was already mentioned, if you really need performance, you probably want to swap std::complex for something else.

I was able to get a decent output for your specific case with all the required shuffles by manually filling small re and im arrays. See: https://godbolt.org/z/sWAAXo This could be trivially extended for ymm registers.

Anyway, here is the ultimate solution adapted from this SO answer which uses intrinsics in combination with clever compiler optimizations:

#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);
    }
}

which compiles to simple

_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