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

Posted 2023-02-16

【中文标题】使用 AVX2 高效计算 std::complex<float> 向量的绝对值

【英文标题】：Efficiently compute absolute values of std::complex<float> vector with AVX2

【发布时间】：2018-12-03 09:53:56

【问题描述】：

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

【问题讨论】：

abs(complex) 是幅度，与二维向量长度相同，sqrt(real^2 + imag^2) (mathwarehouse.com/algebra/complex-number/…)。如果您可以使用 real[i] 和 imag[i] 的单独数组来布置数据，那么它会更有效地进行 SIMD 分析，而无需任何改组，原因与 XY 方向向量相同。

如果你可以使用abs_squared，省略平方根是效率的一大收获。 （与一个乘法 + 1 FMA 相比，sqrt 非常慢。）或者，rsqrtps 和牛顿迭代而不是sqrt 将有助于大多数 CPU 的吞吐量，如果这是您在循环中所做的所有事情 (Fast vectorized rsqrt and reciprocal with SSE/AVX depending on precision) ，但是如果您可以在使用 abs 结果时即时执行此操作，这将是避免 FP sqrt 单元上的吞吐量瓶颈的另一种方法。

谢谢，我知道如何计算它，我也知道不同的数据布局会有所帮助，但是我从无法修改的 API 获取源数据，所以我需要坚持std::complex 的洗牌布局

如果您可以接受 4% 的错误，您还可以查看 alpha-max+beta-min algorithm -- 这也可以使用 SIMD 很好地矢量化。尤其是，如果您需要对输入进行洗牌，那么您可能不会因此而受益匪浅。

@StephenCanon 你是对的，alpha-max-beta-min 通常甚至需要两个额外的and-操作（不过，如果我没看错的话，它们在不同的端口上）。因此，它可能永远不值得在现代架构上实现，除非您遇到下溢/溢出问题（或者，正如您所说，您已经知道 max 和 min）

【参考方案1】：

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

首先将其更改为使用更宽的 256 位寄存器，然后用 __attribute__((aligned (32))) 标记临时 re 和 im 数组以便能够使用对齐加载

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) 高效编译。

【讨论】：

您应该能够使用vpermpd ymm,ymm, imm8 立即进行 64 位随机播放，而不需要在向量中使用随机播放控制。您可能需要一些强制转换内在函数来实现这一点（到 __m256d 并返回）。它可能不会在 Intel CPU 上产生速度差异，但在 Ryzen 上 vpermps 的吞吐量是 vpermpd 的一半。 (agner.org/optimize)。所以那里并没有更糟，并且保存了一个向量寄存器+加载一个向量常量。

您能否详细说明一下 shuffle 是如何按照您提议的方式实现的？

您正在使用 _mm256_permutevar8x32_ps 重新排序 32 位元素对。您可以使用在 64 位元素上运行的 shuffle 执行完全相同的操作。 _mm256_permute4x64_pd( _mm256_castps_pd(abs), _MM_SET(3,1,2,0))，并在此附近返回 __m256。所以它会直接编译为vpermpd 而不是vpermps。

好的，知道了，我用_mm256p_permute_pd试了一下，没用，但这确实是一个不错的解决方案

英特尔的内部搜索页面可以匹配到 asm 助记词。所以你可以在software.intel.com/sites/landingpage/IntrinsicsGuide/… 中输入vpermpd 并找到它。 （是的，由于它们的演化顺序，英特尔的内在函数的名称令人讨厌，令人困惑。Plain permute 是 AVX1 通道内置换，vpermilpd。asm 助记符具有足够的前瞻性，已经描述了它的通道内行为. 这就是我主要记住 asm 助记符并在需要时查找内在函数的原因之一。）

【参考方案2】：

编写复杂 abs 的“高度优化的 AVX2”版本真的很难（如果可能的话），因为标准中定义复数的方式阻止了（特别是由于所有 inf/nan 极端情况）很多优化。

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

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

通过手动填充小的 re 和 im 数组，我能够为您的特定情况获得一个不错的输出，其中包含所有必需的随机播放。见：https://godbolt.org/z/sWAAXo 对于ymm 寄存器，这可以简单地扩展。

无论如何，这里是改编自 this SO answer 的终极解决方案，它使用内在函数与巧妙的编译器优化相结合：

#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

【讨论】：

gcc 将其内联到 mul / fma + vsqrtps，但不会自动矢量化。即使使用-ffast-math，clang 也不会内联。 godbolt.org/z/TBPYd2。即使添加 __restrict 也无助于自动矢量化。我猜 gcc 和 clang 不能胜任洗牌的任务。

@PeterCordes 是的。更改布局后，我能够从 clang 中获得矢量化输出：godbolt.org/z/naseXk

是的，这不足为奇，因为很容易自动矢量化那些可以将 1:1 映射到 asm 的简单垂直操作，但不幸的是，OP 只对交错布局的答案感兴趣。跨度>

@PeterCordes 好吧，在这种情况下，他无能为力，除了添加一些聚集/洗牌和对所有内容进行基准测试以确保这种矢量化有好处。也许，甚至先将给定的std::complex 转换为不同的布局，然后在其上计算 abs 会更有效。

这正是这个问题所要求的：手动矢量化循环。 sqrt 足够昂贵，我相信使用随机播放进行矢量化会有好处。 vsqrtps xmm 在 Intel 和 AMD 上的性能与 sqrtss 相同，但 SQRT 操作的数量是其 4 倍。可能用 vpermps + vextracti128 得到两个 128 位的 real / imag 向量，或者 2x vpermps + 2x vperm2f128 得到两个 256 位的向量。