AVX2:将 4 个复数值与 4 个双精度值相乘和相加的最佳方法是什么?

AVX2: What is the best way to multiply and sum 4 complex values with 4 double values?

xnec2c 项目中 EM 模拟的一个大热点采用这种形式,并且在整个计算过程中以各种方式重复相同的形式:

*dst += (*a1) * (*a2) + (*b1) * (*b2) + (*c1) * (*c2);  // + (*d1) * (*d2)

(This Github link is a real example from matrix.c):

通常 a1、b1、c1 是复数双打。有时 a2、b2、c2 是复数,有时它们是真正的双打。 (我评论了 d1、d2,因为虽然它们在 EM 模拟中不存在,但我猜我们在 AVX 计算中免费获得它们,所以下面的代码写了 4 个双精度数)。

下面是到:

  1. 如上将 4 个复数乘以 4 个双精度数
  2. 求和结果。
  3. 将#2 的总和添加到 *dst
static inline void avx_fma_4zd4d_sum(double complex * restrict dst,
  const double complex * restrict A, 
  double *b0, double *b1, double *b2, double *b3)
{
    // low element first, little-endian style
    __m256d A0 = _mm256_loadu_pd((double *) A); // [A0r A0i A1r A1i ] // [a b c d ]
    __m256d A2 = _mm256_loadu_pd((double *) (A + 2));   // [e f g h ]

    // Q: b0-b3 are spread out in memory. Is this the best way to load them?
    __m256d B0 = _mm256_set_pd(*b1, *b1, *b0, *b0); // reverse from realA order bacause set_pd is.
    __m256d B2 = _mm256_set_pd(*b3, *b3, *b2, *b2); 

    // desired: real=rArB, imag=rBiA
    A0 = _mm256_mul_pd(A0, B0);
    A2 = _mm256_mul_pd(A2, B2);

    // sum A0-A3
    A0 = _mm256_add_pd(A0, A2);

    // Q: Now A0 has 2x complex doubles in a single vector.  How do I sum them?

    // Q: Once I do, how do I add the result to dst?

    //_mm256_storeu_pd((double *) dst, A0);
    //_mm256_storeu_pd((double *) (dst + 2), A2);
}

你可以在上面的评论中看到我的问题。我在看 但它对所有双打的矢量行求和,我的是复数双打所以它看起来像 [A0r, A0i, A1r, A1i] 并且总和需要交错。如果我跳出内在函数,结果会是这样,但我想了解此操作的内在函数细节:

dst += (A0[0]+A0[2]) + (A0[1]+A0[3]) * I;

请注意 *dst 不在 *A 附近的内存中,因此需要单独加载。

我没有 FMA 硬件,但我希望它能够编译成 gcc/clang 如果可用的话会减少到 FMA。

相关问题:

  1. 如果 a2, b2, c2, d2 是复杂的怎么办? (在向量函数中命名为 b0-b3)
    • 当它们都是复数值时,求和到 dst 更容易吗?
  2. 我通常可以将 a1, b1, c1 放在同一个向量中,但是 a2, b2, c2 都在内存中。有没有比使用 _mm256_set_pd 更好的方法来加载它们?

感谢您的帮助!

这是您需要的一些作品。未经测试。

// 4 FP64 complex numbers in 2 AVX SIMD vectors
struct Complex4 { __m256d vec1, vec2; };

// Change the macro if you have FMA and optimizing for Intel
#define USE_FMA 0

Complex4 add( Complex4 a, Complex4 b )
{
    Complex4 res;
    res.vec1 = _mm256_add_pd( a.vec1, b.vec1 );
    res.vec2 = _mm256_add_pd( a.vec2, b.vec2 );
    return res;
}

// Multiply vectors component-wise.
// This is not a complex multiplication, just 2 vmulpd instructions
Complex4 mul_cwise( Complex4 a, Complex4 b )
{
    __m256d r0, r1;
    Complex4 res;
    // Compute the product
    res.vec1 = _mm256_mul_pd( a.vec1, b.vec1 );
    res.vec2 = _mm256_mul_pd( a.vec2, b.vec2 );
    return res;
}

// Multiply + accumulate vectors vectors component-wise
// This is not a complex multiplication.
Complex4 fma_cwise( Complex4 a, Complex4 b, Complex4 acc )
{
#if USE_FMA
    Complex4 res;
    res.vec1 = _mm256_fmadd_pd( a.vec1, b.vec1, acc.vec1 );
    res.vec2 = _mm256_fmadd_pd( a.vec2, b.vec2, acc.vec2 );
    return res;
#else
    return add( mul_cwise( a, b ), acc );
#endif
}

// Load 4 scalars from memory, and duplicate them into 2 AVX vectors
// This is for the complex * real use case
Complex4 load_duplicating( double* rsi )
{
    Complex4 res;
#ifdef __AVX2__
    __m256d tmp = _mm256_loadu_pd( rsi );
    res.vec1 = _mm256_permute4x64_pd( tmp, _MM_SHUFFLE( 1, 1, 0, 0 ) );
    res.vec2 = _mm256_permute4x64_pd( tmp, _MM_SHUFFLE( 3, 3, 2, 2 ) );
#else
    res.vec1 = _mm256_setr_m128d( _mm_loaddup_pd( rsi ), _mm_loaddup_pd( rsi + 1 ) );
    res.vec2 = _mm256_setr_m128d( _mm_loaddup_pd( rsi + 2 ), _mm_loaddup_pd( rsi + 3 ) );
#endif
    return res;
}

// Load 4 complex numbers
Complex4 load_c4( double* rsi )
{
    Complex4 res;
    res.vec1 = _mm256_loadu_pd( rsi );
    res.vec2 = _mm256_loadu_pd( rsi + 4 );
    return res;
}

// Multiply 2 complex numbers by another 2 complex numbers
__m256d mul_CxC_2( __m256d a, __m256d b )
{
    // The formula is [ a.x*b.x - a.y*b.y,  a.x*b.y + a.y*b.x ]
    // Computing it as a.xy * b.xx -+ a.yx * b.yy

    __m256d ayx = _mm256_permute_pd( a, _MM_SHUFFLE2( 0, 1 ) ); // a.yx
    __m256d byy = _mm256_permute_pd( b, _MM_SHUFFLE2( 1, 1 ) ); // b.yy
    __m256d p1 = _mm256_mul_pd( ayx, byy ); // a.yx * b.yy

    __m256d bxx = _mm256_permute_pd( b, _MM_SHUFFLE2( 0, 0 ) ); // b.xx

#if USE_FMA
    return _mm256_fmaddsub_pd( a, bxx, p1 );
#else
    return _mm256_addsub_pd( _mm256_mul_pd( a, bxx ), p1 );
#endif
}

更多笔记。

考虑使用 C++ 编写此类代码。一方面,运算符和重载函数有帮助。此外,您可以使用 Eigen 库用几行代码编写整个热点,其性能通常可与手动编写的内在函数相媲美。

使用编译器标志或 GCC-specific 函数属性确保所有这些函数始终内联。

  1. Add the sum from #2 to *dst

为这样的代码存储中间值从来都不是一个好主意。如果你有 6 个输入 multiply/add,计算它们,并且只存储一次结果。内存访问通常效率低下,read-after-write 尤其如此。

such that gcc/clang will reduce to FMA if available.

将事物简化为 FMA 通常是 Intel 的好主意,不一定是 AMD。 AMD 芯片的 50% 加法/50% 乘法指令组合的吞吐量是 100% FMA 的两倍。与 Intel 不同,在 AMD 上,浮点加法和乘法不会竞争执行单元。 “吞吐量”是指 instructions/second;营销人员决定1次FMA运算=2次浮点运算,所以FLOPs数是一样的。