在AVX通道中混洗的最佳方式

best way to shuffle across AVX lanes?

本文关键字：最佳 方式 混洗 AVX 通道      更新时间：2023-10-16

有一些标题相似的问题，但我的问题涉及一个其他地方没有涉及的非常具体的用例。

我有4个__128d寄存器(x0，x1，x2，x3)，我想在5个__256d寄存器(y0，y1，y2，y3，y4)中重新组合它们的内容，如下所示，为其他计算做准备：

on entry:
x0 contains {a0, a1}
x1 contains {a2, a3}
x2 contains {a4, a5}
x3 contains {a6, a7}
on exit:
y0 contains {a0, a1, a2, a3}
y1 contains {a1, a2, a3, a4}
y2 contains {a2, a3, a4, a5}
y3 contains {a3, a4, a5, a6}
y4 contains {a4, a5, a6, a7}

我在下面的实现相当缓慢。有更好的方法吗？

y0 = _mm256_set_m128d(x1, x0);
__m128d lo = _mm_shuffle_pd(x0, x1, 1);
__m128d hi = _mm_shuffle_pd(x1, x2, 1);
y1 = _mm256_set_m128d(hi, lo);
y2 = _mm256_set_m128d(x2, x1);
lo = hi;
hi = _mm_shuffle_pd(x2, x3, 1);
y3 = _mm256_set_m128d(hi, lo);
y4 = _mm256_set_m128d(x3, x2);

有了寄存器中的输入，您可以在5个混洗指令中完成：

  • 3xvinsertf128，通过串联每个2个xmm寄存器来创建y0、y2和y4
  • 2xvshufpd(通道内混洗)，以创建y1和y3

请注意，y0和y2的低车道包含a1和a2，这是y1低车道所需的元素。同样的洗牌也适用于高车道。

#include <immintrin.h>
void merge(__m128d x0, __m128d x1, __m128d x2, __m128d x3,
__m256d *__restrict y0, __m256d *__restrict y1,
__m256d *__restrict y2, __m256d *__restrict y3, __m256d *__restrict y4)
{
*y0 = _mm256_set_m128d(x1, x0);
*y2 = _mm256_set_m128d(x2, x1);
*y4 = _mm256_set_m128d(x3, x2);
// take the high element from the first vector, low element from the 2nd.
*y1 = _mm256_shuffle_pd(*y0, *y2, 0b0101);
*y3 = _mm256_shuffle_pd(*y2, *y4, 0b0101);
}

非常好地编译(在Godbolt上使用gcc和clang-O3 -march=haswell)为：

merge(double __vector(2), double __vector(2), double __vector(2), double __vector(2), double __vector(4)*, double __vector(4)*, double __vector(4)*, double __vector(4)*, double __vector(4)*):
vinsertf128     ymm0, ymm0, xmm1, 0x1
vinsertf128     ymm3, ymm2, xmm3, 0x1
vinsertf128     ymm1, ymm1, xmm2, 0x1
# vmovapd YMMWORD PTR [rdi], ymm0
vshufpd ymm0, ymm0, ymm1, 5
# vmovapd YMMWORD PTR [rdx], ymm1
vshufpd ymm1, ymm1, ymm3, 5
# vmovapd YMMWORD PTR [r8], ymm3
# vmovapd YMMWORD PTR [rsi], ymm0
# vmovapd YMMWORD PTR [rcx], ymm1
# vzeroupper
# ret

我注释掉了内联中会消失的存储区和东西，所以我们实际上只有5条shuffle指令，而不是问题中的代码的9条shuffler指令。(也包含在Godbolt编译器资源管理器链接中)。

这在AMD上非常好，因为vinsertf128非常便宜(因为256位寄存器被实现为2x128位的一半，所以它只是一个128位的副本，不需要特殊的混洗端口。)256位通道交叉混洗在AMD上很慢，但像vshufpd这样的通道内256位混洗只有2个uops。

在英特尔上，这很好，但主流的带有AVX的英特尔CPU对于256位或FP混洗，每个时钟只有1个混洗吞吐量。(Sandybridge和更早的版本对整数128位混洗有更高的吞吐量，但AVX2 CPU放弃了额外的混洗单元，它们对此毫无帮助。)

因此，英特尔CPU根本无法利用指令级并行性，但总共只有5个uops，这很好。这是最低限度的可能，因为你需要5个结果。

但特别是如果周围的代码也在混洗时遇到瓶颈，值得考虑一种只有4个存储和5个重叠矢量加载的存储/重新加载策略。或者可能是2xvinsertf128来构造y0和y4，然后2x 256位存储+3个重叠的重新加载。这可能会让无序的exec仅使用y0或y4开始执行相关指令，而存储转发暂停已为y1..3解决。

特别是如果你不太关心英特尔第一代Sandybridge，在那里，未对齐的256位矢量加载效率较低。(请注意，如果您使用GCC，您可能希望使用gcc -mtune=haswell进行编译，以关闭-mavx256-split-unaligned-load默认/sandybridge调优。无论编译器是什么，如果使二进制文件在编译它的机器上运行，以充分利用指令集和设置调优选项，-march=native都是一个好主意。)

但是，如果前端的总uop吞吐量在瓶颈所在的地方更大，那么shuffle实现是最好的。

(请参见https://agner.org/optimize/以及x86标签wiki中的其他性能链接，以了解有关性能调优的更多信息。此外，在预测现代超标量处理器上的操作延迟时需要考虑哪些因素，以及我如何手动计算它们？，但实际上，Agner Fog的指南是一个更深入的指南，解释了吞吐量与延迟的实际关系。)

我甚至不需要保存，因为数据在连续内存中也已经可用。

然后简单地用5个重叠的负载加载它几乎可以肯定是你能做的最有效的事情。

Haswell可以在L1d的每个时钟上进行2次加载，或者在任何一次跨越缓存线边界时都可以进行更少的加载因此，如果您可以将块对齐64，则完全没有缓存线分割，非常高效缓存未命中很慢，但从L1d缓存重新加载热数据非常便宜，而且支持AVX的现代CPU通常具有高效的未对齐加载支持。

(正如我之前所说，如果使用gcc，请确保使用-march=haswell或-mtune=haswell进行编译，而不仅仅是-mavx，以避免gcc的-mavx256-split-unaligned-load。)

4个负载+1个vshufpd(y0，y2)可能是平衡负载端口压力和ALU压力的好方法，具体取决于周围代码中的瓶颈。或者，如果周围代码的混洗端口压力较低，甚至可以进行3次加载+2次混洗。

它们位于需要加载它们的先前计算的寄存器中。

如果之前的计算在寄存器中仍然有源数据，那么您可以首先加载256位，并将其128位的下半部分用于之前的计算。(XMM寄存器是相应YMM寄存器的低128，读取它们不会干扰上层通道，因此_mm256_castpd256_pd128编译为零asm指令。)

对y0、y2和y4执行256位加载，并将它们的下半部分用作x0、x1和x2。(稍后使用未对齐的加载或搅乱来构造y1和y3)。

只有x3还不是您想要的256位矢量的低128位。

理想情况下，当您从同一地址执行_mm_loadu_pd和_mm256_loadu_pd时，编译器已经注意到了这种优化，但您可能需要通过执行来手动保持它

__m256d y0 = _mm256_loadu_pd(base);
__m128d x0 = _mm256_castpd256_pd128(y0);

等等，以及提取ALU内在(_mm256_extractf128_pd)或x3的128位负载，这取决于周围的代码。如果它只需要一次，那么让它折叠成内存操作数以供任何指令使用可能是最好的。

潜在的缺点：在128位计算开始之前的延迟稍高，或者如果256位加载是缓存线交叉，而128位加载不是，则需要几个周期。但是，如果您的数据块按64字节对齐，则不会发生这种情况。

相关文章：