如何使用256位AVX向量平方两个复数双精度数？

Question

Matt Scarpino详细解释了如何使用英特尔的AVX指令集来乘以两个复数双精度数（虽然他承认这可能不是最优算法，但我非常感激他）。以下是他的方法，我已经验证过：

__m256d vec1 = _mm256_setr_pd(4.0, 5.0, 13.0, 6.0);
__m256d vec2 = _mm256_setr_pd(9.0, 3.0, 6.0, 7.0);
__m256d neg  = _mm256_setr_pd(1.0, -1.0, 1.0, -1.0);

/* Step 1: Multiply vec1 and vec2 */
__m256d vec3 = _mm256_mul_pd(vec1, vec2);

/* Step 2: Switch the real and imaginary elements of vec2 */
vec2 = _mm256_permute_pd(vec2, 0x5);

/* Step 3: Negate the imaginary elements of vec2 */
vec2 = _mm256_mul_pd(vec2, neg);  

/* Step 4: Multiply vec1 and the modified vec2 */
__m256d vec4 = _mm256_mul_pd(vec1, vec2);

/* Horizontally subtract the elements in vec3 and vec4 */
vec1 = _mm256_hsub_pd(vec3, vec4);

/* Display the elements of the result vector */
double* res = (double*)&vec1;
printf("%lf %lf %lf %lf\n", res[0], res[1], res[2], res[3]);

我的问题是我想计算两个复数的平方。我尝试了使用马特的技巧如下：

struct cmplx a;
struct cmplx b;

a.r = 2.5341;
a.i = 1.843;

b.r = 1.3941;
b.i = 0.93;

__m256d zzs = squareZ(a, b);

double* res = (double*) &zzs;

printf("\nA: %f + %f,  B: %f + %f\n", res[0], res[1], res[2], res[3]);

使用 Haskell 的复数算术，我已经验证了结果的正确性，除了，如您所见，B 的实部：

A: 3.025014 + 9.340693,  B: 0.000000 + 2.593026

我有两个问题：是否有更好的（更简单和/或更快）方法使用AVX指令来平方两个复数双精度数？如果没有，我该如何修改Matt的代码来实现呢？

- sacheie

---

squareZ是什么？ - mtrw

请使用“编辑”按钮更新代码。当前的情况是，没有人可以剪切和粘贴您的代码并运行它来重现问题。请参见 http://stackoverflow.com/help/mcve。 - mtrw

对奇数元素（保持虚部）取反最好使用 _mm256_xor_pd 来翻转符号位，这是 FP 版本向量 XOR/AND/OR 的主要用例之一。同时需要注意的是将实部和虚部放在两个单独的向量中可以实现更高效的操作：1. 不需要洗牌去交换实部/虚部：通过操作不同的变量就可以免费完成。2. 快速垂直子，而不是慢速水平子。3. 每条指令可以翻转 4 个虚部的符号。使用 shuffle_pd 对打包输入/输出执行此操作可能会获得胜利。 - Peter Cordes

是的，我看到太多人坚持使用旧的实部/虚部交替表示方式，这让我很生气。在很多情况下，糟糕的表示方式是由API层面或为了向后兼容性而规定的。然后人们抱怨编译器无法像您和我这样矢量化一个std::complex数组... - Mysticial

@PeterCordes 我相信 std::complex 遵循一些 IEEE 规范来处理通常的 NaN 边界情况 - 这扩展到了 这个怪物。这甚至不好笑。 - Mysticial

Answer 1

本答案涵盖了复数数组乘法的一般情况

理想情况下，将数据存储在单独的实数和虚数数组中，这样您就可以加载连续的实部和虚部向量。这使得交叉相乘变得自由（只需使用不同的寄存器/变量），而无需在向量内部重新排列内容。

您可以在飞行中以相对便宜的方式在交错的double complex样式和SIMD友好的单独向量样式之间进行转换，但受AVX在通道内洗牌的影响。例如，使用unpacklo / unpackhi shuffle相当便宜，可以在一个通道内解交错或重新交错，如果您不关心临时向量内数据的实际顺序。

实际上，这种洗牌非常便宜，即使针对单个复杂乘积，在支持FMA的CPU上也比Matt的代码（即使是经过调整的版本）更为出色。这需要以4个复杂双精度（2个结果向量）的组形式产生结果。

如果您只需要一次生成一个结果向量，则我还提出了一种替代Matt算法的方法，可以使用FMA（实际上是FMADDSUB）并避免单独的符号更改指令。

---

如果你使用了-ffast-math，gcc会自动将简单的复数乘法标量循环向量化成相当不错的代码。它会像我建议的那样进行解开交错。

#include <complex.h>

// even with -ffast-math -ffp-contract=fast, clang doesn't manage to use vfmaddsubpd, instead using vmulpd and vaddsubpd :(
// gcc does use FMA though.

// auto-vectorizes with a lot of extra shuffles
void cmul(double complex *restrict dst,
          const double complex *restrict A, const double complex *restrict B)
{       // clang and gcc change strategy slightly for i<1 or i<2, vs. i<4
  for (int i=0; i<4 ; i++) {
    dst[i] = A[i] * B[i];
  }
}

请看Godbolt编译器资源管理器上的汇编代码。我不确定clang的汇编质量如何；它使用了许多64b->128bVMODDDUP广播加载。这种形式在Intel CPU上纯粹在加载端口中处理（请参见Agner Fog's insn tables），但仍需要执行很多操作。如前所述，gcc使用4个VPERMPD洗牌来重新排列各行，在乘法/ FMA之前，然后使用另外4个VPERMPD将结果重新排序，再与VSHUFPD组合。这是4个复杂乘法需要额外进行8次洗牌。

---

将gcc版本转换回内在函数并去除冗余的洗牌操作可以得到最优代码。（显然，gcc希望其临时变量按照A B C D顺序排列，而不是由VUNPCKLPD (_mm256_unpacklo_pd)的in-lane行为导致的A C B D顺序。）

我将代码on Godbolt与Matt的代码进行了微调，让你可以尝试不同的编译器选项和版本。

// multiplies 4 complex doubles each from A and B, storing the result in dst[0..3]
void cmul_manualvec(double complex *restrict dst,
          const double complex *restrict A, const double complex *restrict B)
{
                                         // 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 ]
  __m256d realA = _mm256_unpacklo_pd(A0, A2);  // [A0r A2r  A1r A3r ] // [a e c g ]
  __m256d imagA = _mm256_unpackhi_pd(A0, A2);  // [A0i A2i  A1i A3i ] // [b f d h ]
  // the in-lane behaviour of this interleaving is matched by the same in-lane behaviour when we recombine.
  
  __m256d B0 = _mm256_loadu_pd((double*)B);                           // [m n o p]
  __m256d B2 = _mm256_loadu_pd((double*)(B+2));                       // [q r s t]
  __m256d realB = _mm256_unpacklo_pd(B0, B2);                         // [m q o s]
  __m256d imagB = _mm256_unpackhi_pd(B0, B2);                         // [n r p t]

  // desired:  real=rArB - iAiB,  imag=rAiB + rBiA
  __m256d realprod = _mm256_mul_pd(realA, realB);
  __m256d imagprod = _mm256_mul_pd(imagA, imagB);
  
  __m256d rAiB     = _mm256_mul_pd(realA, imagB);
  __m256d rBiA     = _mm256_mul_pd(realB, imagA);

  // gcc and clang will contract these nto FMA.  (clang needs -ffp-contract=fast)
  // Doing it manually would remove the option to compile for non-FMA targets
  __m256d real     = _mm256_sub_pd(realprod, imagprod);  // [D0r D2r | D1r D3r]
  __m256d imag     = _mm256_add_pd(rAiB, rBiA);          // [D0i D2i | D1i D3i]

  // interleave the separate real and imaginary vectors back into packed format
  __m256d dst0 = _mm256_shuffle_pd(real, imag, 0b0000);  // [D0r D0i | D1r D1i]
  __m256d dst2 = _mm256_shuffle_pd(real, imag, 0b1111);  // [D2r D2i | D3r D3i]
  _mm256_storeu_pd((double*)dst, dst0);
  _mm256_storeu_pd((double*)(dst+2), dst2);   
}

  Godbolt asm output: gcc6.2 -O3 -ffast-math -ffp-contract=fast -march=haswell
    vmovupd         ymm0, YMMWORD PTR [rsi+32]
    vmovupd         ymm3, YMMWORD PTR [rsi]
    vmovupd         ymm1, YMMWORD PTR [rdx]
    vunpcklpd       ymm5, ymm3, ymm0
    vunpckhpd       ymm3, ymm3, ymm0
    vmovupd         ymm0, YMMWORD PTR [rdx+32]
    vunpcklpd       ymm4, ymm1, ymm0
    vunpckhpd       ymm1, ymm1, ymm0
    vmulpd          ymm2, ymm1, ymm3
    vmulpd          ymm0, ymm4, ymm3
    vfmsub231pd     ymm2, ymm4, ymm5     # separate mul/sub contracted into FMA
    vfmadd231pd     ymm0, ymm1, ymm5
    vunpcklpd       ymm1, ymm2, ymm0
    vunpckhpd       ymm0, ymm2, ymm0
    vmovupd         YMMWORD PTR [rdi], ymm1
    vmovupd         YMMWORD PTR [rdi+32], ymm0
    vzeroupper
    ret

---

对于4个复杂数乘法（2对输入向量），我的代码使用：

• 4次加载（每个32B）
• 2次存储（每个32B）
• 6个in-lane洗牌（每个输入向量一个，每个输出一个）
• 2个VMULPD
• 2个VFMA...something
• （如果我们可以使用分离的实部和虚部向量或者输入已经处于这种格式，则只需4个洗牌；否则不需洗牌）
• 在Intel Skylake上的延迟（不包括加载/存储）：14个周期= 4个洗牌周期直到第二个VMULPD开始+ 4个周期（第二个VMULPD）+ 4c（第二个vfmadd231pd）+ 1c（第一个结果向量准备好的1c洗牌）+ 1c（第二个结果向量洗牌）

因此，在吞吐量方面，这完全受到洗牌端口的瓶颈。（每个时钟周期1个洗牌吞吐量，与Intel Haswell及更高版本的每个时钟周期总共2个MUL / FMA / ADD相比）。这就是为什么打包存储很糟糕的原因：洗牌的吞吐量有限，花费更多的指令来洗牌而不是做数学运算并不好。

---

Matt Scarpino的代码，我做了一些小改动（重复4次以执行4个复杂的乘法）。 （有关更有效地生成一个向量的我的重写，请参见下文）。

• 相同的6个加载/存储操作
• 6个通道内洗牌操作（在当前的英特尔和AMD CPU上，HSUBPD是2个洗牌操作和一个减法操作）
• 4个乘法操作
• 2个减法操作（不能与乘法合并为FMA操作）
• 额外的一条指令（加上一个常数）来翻转虚数元素的符号。Matt选择了乘以1.0或-1.0，但高效的选择是使用异或位运算（即XORPD与-0.0）。
• 在Intel Skylake上第一个结果向量的延迟为11个周期。1个周期（vpermilpd和vxorpd在同一周期内）+ 4个周期（第二个vmulpd）+ 6个周期（vhsubpd）。第一个vmulpd与其他操作重叠，在洗牌和vxorpd的同一周期开始。计算第二个结果向量应该相当不错地交错进行。

马特代码的主要优点是一次只需使用一个向量宽度的复数乘法，而不需要你拥有4个输入数据向量。它具有较低的延迟。但请注意，我的版本不需要2对输入向量来自连续内存，也不需要彼此相关。它们在处理过程中混合在一起，但结果是2个独立的32B向量。

我的修改版马特代码在没有FMA的CPU上几乎和4个同时进行的版本一样好（只需要额外的VXORPD），但当阻止我们利用FMA时，效果显著更差。此外，它永远不会以非打包形式提供结果，因此您无法将分离的形式用作另一个乘法的输入并跳过洗牌。

---

使用FMA逐个向量计算结果:

如果有时候需要平方而不是乘以两个不同的复数，请不要使用此方法。这类似于Matt的算法，公共子表达式消除并没有简化。

我还没有键入此方法的C内部函数，只是计算了数据移动。由于所有的洗牌都在一个通道内，我只会展示低通道。使用相关指令的256b版本可以在两个通道中执行相同的洗牌。它们保持分开。

// MULTIPLY: for each AVX lane: am-bn, an+bm  
 r i  r i
 a b  c d       // input vectors: a + b*i, etc.
 m n  o p

算法:

• 使用 movshdup(a b) + mulpd 创建 bm bn

• 在前一个结果上使用 shufpd 创建 bn bm。（或者在乘法之前使用 shuffle 创建 n m）

• 使用 movsldup(a b) 创建 a a

• 使用 fmaddsubpd 生成最终结果：[a|a]*[m|n] -/+ [bn|bm]。

  是的，SSE/AVX 有 ADDSUBPD 来在偶数/奇数元素中交替进行减法/加法（按照这种用例的顺序）。FMA 包括 FMADDSUB132PD，它执行减法和加法（以及反向操作 FMSUBADD，执行加法和减法）。

每4个结果：6倍洗牌，2倍乘法，2倍FMADD/SUB。 因此，除非我错了，它与去交错的方法一样有效（当不平方相同的数字时）。 Skylake延迟= 10c =1+4+1来创建bn bm（与创建a a重叠1个周期），+ 4（FMA）。 因此，它比Matt的延迟低一个周期。

在Bulldozer系列上，将两个输入都洗牌到第一个乘法中会是胜利，因此乘法->fmaddsub关键路径保持在FMA域内（延迟降低1个周期）。 另一种做法可防止愚蠢的编译器过早进行movsldup（a b）并延迟mulpd。 （在循环中，许多迭代将正在飞行并且将在洗牌端口瓶颈。）

---

这个方法相对于Matt的方法来说在平方运算上仍然更好（仍然可以使用XOR并且可以使用FMA），但我们无法节省任何洗牌操作。

// SQUARING: for each AVX lane: aa-bb, 2*ab
// ab bb  // movshdup + mul
// bb ab  // ^ -> shufpd

// a  a          // movsldup
// aa-bb  ab+ab  // fmaddsubpd : [a|a]*[a|b] -/+ [bb|ab]
// per 4 results: 6x shuffle, 2x mul, 2xfmaddsub

---

我也尝试了一些可能性，比如 (a+b) * (a+b) = aa+2ab+bb，或者 (r-i)*(r+i) = rr - ii，但是没有任何进展。在步骤之间进行四舍五入意味着浮点数计算无法完全抵消，因此做这样的事情甚至不会产生完全相同的结果。

- Peter Cordes

---

哇，这个答案非常详细，我希望我可以投多次赞。我曾考虑过将实部和虚部分别存储在不同的数组中以获得更好的效率。但是我认为可能会失去一些东西（仅仅因为很多信息超出了我的理解范围），我希望问题能够简化，因为这两个复数只需要被平方。我不需要两个任意复双精度数的乘积；我只需要尽可能高效地计算它们的平方。如果你已经解释过了而我没有注意到，请指出我错过了哪里。 - sacheie

@sacheie 当你阅读彼得的回答时，你应该会被淹没在其中。没有例外，即使是我自己也不例外。 - Mysticial

@Mysticial：如果我在回答问题时没有记录下每一个相关的想法，那会感觉很浪费:P像这样的静态性能分析是比较不同实现短SIMD函数的几种有用方法之一，所以除了在我没有的CPU上进行基准测试外，这是回答“是否有更好的方法”的唯一途径。我很高兴人们发现它有用:) - Peter Cordes

@Mysticial：我还没有尝试过为Agner Fog尚未撰写的任何内容进行调整。虽然我通常不会接触到最新的硬件，但这是类似的情况。我编写了这个使用PEXT动态生成AVX2 + BMI2左侧包和洗牌掩码的代码，除了确保它编译成良好的汇编代码之外，没有进行任何测试。OP报告说它可以正常工作：）。我认为在Haswell上它已经接近最优了，但我还没有真正尝试过。 - Peter Cordes

啊，是啊，我的程序员一面加上超频者一面就太强了。所以我最终会得到每一个世代的一部分。 - Mysticial

Answer 2

请参考我的另一篇答案，其中有关于不同复数相乘的一般情况的内容，而非仅限于平方。

简要总结：只需使用我的其他答案中的代码，并将两个输入值设为相同。编译器会处理冗余部分。

---

平方可以稍微简化一下运算：我们只需要计算rAiB和rBiA两个不同的向量积即可。但是最终结果还是需要加倍，因此我们最终需要进行2次乘法、1次FMA和1次加法，而不是2次乘法和2次FMA。

在SIMD不友好的交错存储格式下，这种方法能够大幅提高解缠绕方法的效率，因为只需要对一个输入进行混洗。Matt的方法则没有任何收益，因为它使用相同的向量乘积计算两个向量积。

---

使用我在另一个答案中给出的cmul_manualvec()函数：

// squares 4 complex doubles from A[0..3], storing the result in dst[0..3]
void csquare_manual(double complex *restrict dst,
          const double complex *restrict A) {
  cmul_manualvec(dst, A, A);
}

gcc和clang都足够聪明，可以优化掉两次使用相同输入的冗余，所以没有必要使用内部函数创建自定义版本。clang在标量自动矢量化版本上表现不佳，因此不要使用它。我认为与此asm输出相比没有任何收益

        clang3.9 -O3 -ffast-math -ffp-contract=fast -march=haswell
    vmovupd         ymm0, ymmword ptr [rsi]
    vmovupd         ymm1, ymmword ptr [rsi + 32]
    vunpcklpd       ymm2, ymm0, ymm1
    vunpckhpd       ymm0, ymm0, ymm1   # doing this shuffle first would let the first multiply start a cycle earlier.  Silly compiler.
    vmulpd          ymm1, ymm0, ymm0   # imag*imag
    vfmsub231pd     ymm1, ymm2, ymm2   # real*real - imag*imag
    vaddpd          ymm0, ymm0, ymm0   # imag+imag = 2*imag
    vmulpd          ymm0, ymm2, ymm0   # 2*imag * real
    vunpcklpd       ymm2, ymm1, ymm0
    vunpckhpd       ymm0, ymm1, ymm0
    vmovupd ymmword ptr [rdi], ymm2
    vmovupd ymmword ptr [rdi + 32], ymm0
    vzeroupper
    ret

可能使用不同的指令排序会更好，例如将实向量加倍，因为它首先被展开，所以VADDPD可以在imag*imag VMULPD之前更早地启动一个周期，从而减少资源冲突。但是，在C源代码中重新排列行通常不能直接翻译为汇编重排序，因为现代编译器是复杂的工具。 （据我所知，gcc并不特别尝试为x86调度指令，因为乱序执行大多数情况下隐藏了这些效果。）

无论如何，每4个复平方：

2个加载（从4个降至2个）+ 2个存储，原因显而易见

4个洗牌（从6个降至4个），同样显而易见

2个VMULPD（相同）

1个FMA + 1个VADDPD（从2 FMA降至1个FMA。在Haswell / Broadwell上，VADDPD比FMA具有较低的延迟；在Skylake上也是如此。）

马特的版本仍然是6个洗牌，其他所有内容都相同。

- Peter Cordes

---

我已经完全复制了你的代码（除了将名称更改为cproduct和csquare），结果是完美的，但是我立即遇到了segfault。 - sacheie

以下是我的Makefile标志：COMPILER_FLAGS = -std=c11 -O3 -Wall -Wextra -march=skylake -mtune=skylake -mavx2 - sacheie

使用gcc编译器，并且这是我的代码：int main() { double complex dst; double complex a = 2.5341 + 1.843 * I; csquare(&dst, &a); printf("%f + %f\n", creal(dst), cimag(dst)); return 0; }关于'restrict'关键字，我是否有什么误解？ - sacheie

@sacheie。这并不奇怪。请注意在cproduct的顶部有一个注释// multiplies 4 complex doubles each from A and B, storing the result in dst[0..3]，描述了它读取和写入的内存。至于restrict，如果http://en.cppreference.com/w/c/language/restrict 描述的是你认为的，那么不，否则是的，你误解了它。我不知道你认为它是什么！ - Peter Cordes

啊，我真是太蠢了。谢谢你的耐心等待。我忘记正确引用dst了；这与“restrict”无关。 - sacheie