使用 AVX 对 64 位结构进行排序？

Question

我有一个 64 位结构，它表示几条数据，其中之一是浮点值：

struct MyStruct{
    uint16_t a;
    uint16_t b;
    float f;
}; 

我有四个这样的结构，比如说std::array<MyStruct, 4>

是否可以使用 AVX 对数组进行排序，根据浮点成员 MyStruct::f？

Answer 1

对不起，这个答案很混乱；它并没有一次全部写完，我很懒。有一些重复。

我有 4 个不同的想法：

1. 正常排序，但将结构作为 64 位单元移动
2. 向量化插入排序作为 qsort 的构建块

3. 排序网络，比较器实现使用cmpps / blendvpd 而不是minps/maxps。不过，额外的开销可能会扼杀加速。

4. 对网络进行排序：加载一些结构，然后混洗/混合以获得一些仅浮点数的寄存器和一些仅有效负载的寄存器。使用 Timothy Furtak 的技术做一个普通的minps/maxps 比较器，然后cmpeqps min,orig -> 对有效负载进行屏蔽异或交换。这对每个比较器排序两倍的数据，但确实需要在比较器之间的两个寄存器上匹配洗牌。完成后还需要重新交错（但使用 unpcklps / unpckhps 很容易，如果您安排比较器，以便那些通道内解包将最终数据按正确顺序放置）。

   这也避免了某些 CPU 在对负载中表示非正规、NaN 或无穷大的位模式进行 FP 比较时可能出现的潜在减速，而无需在 MXCSR 中设置非正规为零位。

   Furtak 的论文建议在使用向量进行排序后进行标量清理，这将大大减少洗牌的数量。

正常排序

在使用普通排序算法时，通过使用 64 位加载/存储移动整个结构，并对 FP 元素进行标量 FP 比较，至少可以获得小幅加速。为了使这个想法尽可能地发挥作用，首先使用浮点值对结构进行排序，然后您可以movq 将整个结构放入 xmm reg，并且浮点值将在 low32 中ucomiss。然后你（或者可能是一个智能编译器）可以用movq 存储结构。

看看 Kerrek SB 链接到的 asm 输出，编译器似乎在有效地复制结构方面做得相当糟糕：

icc 似乎分别 movzx 两个 uint 值，而不是在 64b 负载中舀出整个结构。也许它不打包结构？ gcc 5.1 似乎大部分时间都没有这个问题。

加速插入排序

对于足够小的问题，大排序通常与插入排序分而治之。 Insertion sort 将数组元素复制一个，只有当我们发现我们已经到达当前元素所属的位置时才会停止。所以我们需要将一个元素与一系列打包元素进行比较，如果比较为真则停止。你闻到矢量的味道吗？我闻到了矢量的味道。

# RSI points to  struct { float f; uint... payload; } buf[];
# RDI points to the next element to be inserted into the sorted portion
# [ rsi to rdi ) is sorted, the rest isn't.
##### PROOF OF CONCEPT: debug / finish writing before using!  ######

.new_elem:
vbroadcastsd ymm0, [rdi]      # broadcast the whole struct
mov rdx, rdi

.search_loop:
    sub        rdx, 32
    vmovups    ymm1, [rdx]    # load some sorted data
    vcmplt_oqps ymm2, ymm0, ymm1   # all-ones in any element where ymm0[i] < ymm1[i] (FP compare, false if either is NaN).
    vmovups    [rdx+8], ymm1  # shuffle it over to make space, usual insertion-sort style
    cmp        rdx, rsi
    jbe     .endsearch        # below-or-equal (addresses are unsigned)
    movmskps   eax, ymm2
    test       al, 0b01010101 # test only the compare results for 

    jz      .search_loop      # [rdi] wasn't less than any of the 4 elements

.endsearch:
# TODO: scalar loop to find out where the new element goes.
#  All we know is that it's less than one of the elements in ymm1, but not which
add           rdi, 8
vmovsd         [rdx], ymm0
cmp           rdi, r8   # pointer to the end of the buf
jle           .new_elem

  # worse alternative to movmskps / test:
  # vtestps    ymm2, ymm7     # where ymm7 is loaded with 1s in the odd (float) elements, and 0s in the even (payload) elements.
  # vtestps is like PTEST, but only tests the high bit.  If the struct was in the other order, with the float high, vtestpd against a register of all-1s would work, as that's more convenient to generate.

这肯定充满了错误，我应该用 C 语言用内在函数编写它。

这是一种插入排序，其开销可能比大多数类型都多，由于处理前几个元素（不填充向量）和在跳出检查多个元素的向量搜索循环后，找出新元素的放置位置。

可能对循环进行流水线化，因此我们在下一次迭代（或中断之后）之前没有存储ymm1，这样可以节省冗余存储。通过移位/改组它们在寄存器中进行比较，而不是从字面上进行标量加载/比较可能是一个胜利。这可能会导致太多不可预测的分支，而且我没有看到一个很好的方法来结束 vmovups 的 reg 中的高 4 和 vmovsd 的另一个 reg 中的低 4。

我可能发明了一种两全其美的插入排序：对于小数组来说很慢，因为在跳出搜索循环后需要做更多的工作，但它仍然是插入排序：对于大数组来说很慢，因为 O(n^2 ）。但是，如果可以使 searchloop 之外的代码变得不可怕，这可能会作为 qsort / mergesort 的小数组端点有用。

无论如何，如果有人将这个想法开发成实际调试和工作的代码，请告诉我们。

更新：Timothy Furtak's paper 描述了用于对短数组进行排序的 SSE 实现（用作更大排序的构建块，例如这种插入排序）。他建议使用 SSE 生成部分有序的结果，然后使用标量操作进行清理。 （对大多数排序的数组进行插入排序很快。）

这导致我们：

排序网络

这里可能没有任何加速。 Xiaochen、Rocki 和 Suda 仅报告在 Xeon Phi 卡上，对于 32 位 (int) 元素、单线程合并排序，从 scalar -> AVX-512 的速度提升了 3.7 倍。使用更宽的元素，更少适合向量 reg。 （这对我们来说是 4 倍：256b 中有 64b 元素，而 512b 中有 32b 元素。）他们还利用 AVX512 掩码仅比较一些通道，这是 AVX 中不可用的功能。此外，由于比较器功能较慢，竞争 shuffle/blend 单元，我们已经处于更糟糕的状态。

Sorting networks 可以使用 SSE/AVX 压缩比较指令构造。 （更常见的是，使用一对最小/最大指令可以有效地执行一组打包的 2 元素排序。）更大的排序可以通过执行成对排序的操作来构建。 This paper by Tian Xiaochen, Kamil Rocki and Reiji Suda at U of Tokyo 有一些真正的 AVX 代码用于排序（没有有效负载），并讨论了向量寄存器的棘手之处，因为您无法比较同一寄存器中的两个元素（因此必须将排序网络设计为不需要那）。他们使用pshufd 来排列元素以进行下一次比较，从而建立一个更大的排序，而不是仅对几个满是数据的寄存器进行排序。

现在，诀窍是根据仅半个元素的比较来做一种成对的 64b 元素对。 （即使用排序键保留有效负载。）我们可以通过对(key, payload) 对的数组进行排序来对其他事物进行排序，其中有效负载可以是索引或32 位指针（mmap(MAP_32bit) 或x32 ABI）。

所以让我们自己构建一个比较器。用排序网络的说法，这是对一对输入进行排序的操作。所以它要么在寄存器之间交换一个元素，要么不交换。

# AVX comparator for SnB/IvB
# struct { uint16_t a, b; float f; }  inputs in ymm0, ymm1
# NOTE: struct order with f second saves a shuffle to extend the mask

vcmpps    ymm7, ymm0, ymm1, _CMP_LT_OQ  # imm8=17: less-than, ordered, quiet (non-signalling on NaN)
     # ymm7 32bit elements = 0xFFFFFFFF if ymm0[i] < ymm1[i], else 0
# vblendvpd checks the high bit of the 64b element, so mask *doesn't* need to be extended to the low32
vblendvpd ymm2, ymm1, ymm0, ymm7
vblendvpd ymm3, ymm0, ymm1, ymm7
# result: !(ymm2[i] > ymm3[i])  (i.e. ymm2[i] < ymm3[i], or they're equal or unordered (NaN).)
#  UNTESTED

您可能需要设置 MXCSR 以确保 int 位在碰巧表示非正规或 NaN 浮点数时不会减慢您的 FP 操作。我不确定这是否只发生在 mul/div 上，或者它是否会影响比较。

• Intel Haswell：延迟：ymm2 准备好 5 个周期，ymm3 准备好 7 个周期。吞吐量：每 4 个周期一个。 （p5 瓶颈）。
• 英特尔 Sandybridge/Ivybridge：延迟：ymm2 准备就绪需要 5 个周期，ymm3 需要 6 个周期。吞吐量：每 2 个周期一个。 （p0/p5 瓶颈）。
• AMD Bulldozer/Piledriver: (vblendvpd ymm: 2c lat, 2c recip tput): lat: 4c for ymm2, 6c for ymm3。或者更糟糕的是，在 cmpps 和 blend 之间存在旁路延迟。 tput：每 4c 一个。 （向量 P1 的瓶颈）
• AMD Steamroller: (vblendvpd ymm: 2c lat, 1c recip tput): lat: 4c for ymm2, 5c for ymm3。或者可能由于旁路延迟而高出 1 个。 tput：每 3c 一个（矢量端口 P0/1 上的瓶颈，用于 cmp 和 blend）。

VBLENDVPD 是 2 微秒。 （它有 3 个 reg 输入，所以它不能是 1 uop :/）。两个 uops 都只能在 shuffle 端口上运行。在 Haswell 上，这只是端口 5。在 SnB 上，这是 p0/p5。 （IDK 为什么 Haswell 将 shuffle / blend 吞吐量与 SnB/IvB 相比减半。）

如果 AMD 设计有 256b 宽的向量单元，它们的低延迟 FP 比较和 3 输入指令的单宏操作解码将使它们领先。

通常的 minps/maxps 对是 3 和 4 个周期的延迟 (ymm2/3)，以及每 2 个周期的吞吐量 (Intel)。 （FP 添加/子/比较单元上的 p1 瓶颈）。最公平的比较可能是对 64 位双精度数进行排序。如果没有要比较的多对独立寄存器，则额外的延迟可能会受到伤害。 Haswell 上减半的吞吐量将大大降低任何加速。

另外请记住，比较器操作之间需要进行洗牌，以便将正确的元素排列好进行比较。 min/maxps 未使用 shuffle 端口，但我的 cmpps/blendv 版本使它们饱和，这意味着 shuffle 不能与比较重叠，除非是为了填补数据依赖关系留下的空白。

使用超线程，另一个可以保持其他端口忙碌的线程（例如，端口 0/1 fp mul/add 单元或整数代码）将与这个混合瓶颈版本很好地共享一个内核。

我尝试了 Haswell 的另一个版本，它使用按位 AND/OR 运算“手动”进行混合。不过，它最终变慢了，因为在合并之前，两个来源都必须双向屏蔽。

# AVX2 comparator for Haswell
# struct { float f; uint16_t a, b; }  inputs in ymm0, ymm1
#
vcmpps ymm7, ymm0, ymm1, _CMP_LT_OQ  # imm8=17: less-than, ordered, quiet (non-signalling on NaN)
     # ymm7 32bit elements = 0xFFFFFFFF if ymm0[i] < ymm1[i], else 0
vshufps ymm7, ymm7, ymm7, mask(0, 0, 2, 2)  # extend the mask to the payload part.  There's no mask function, I just don't want to work out the result in my head.
vpand    ymm10, ymm7, ymm0       # ymm10 = ymm0 keeping elements where ymm0[i] < ymm1[i]
vpandn   ymm11, ymm7, ymm1       # ymm11 = ymm1 keeping elements where !(ymm0[i] < ymm1[i])
vpor     ymm2, ymm10, ymm11      # ymm2 = min_packed_mystruct(ymm0, ymm1)

vpandn   ymm10, ymm7, ymm0       # ymm10 = ymm0 keeping elements where !(ymm0[i] < ymm1[i])
vpand    ymm11, ymm7, ymm1       # ymm11 = ymm1 keeping elements where ymm0[i] < ymm1[i]
vpor     ymm3, ymm10, ymm11  # ymm2 = max_packed_mystruct(ymm0, ymm1)

# result: !(ymm2[i] > ymm3[i])
#  UNTESTED

这是 8 微指令，而 blendv 版本是 5。最后 6 条和/和/或指令中有很多并行性。不过，cmpps 有 3 个周期延迟。我认为ymm2 将在 6 个周期内准备好，而 ymm3 在 7 个周期内准备好。（并且可以与 ymm2 上的操作重叠）。比较器操作之后的 insns 可能会被洗牌，以便将数据放入正确的元素中以进行下一次比较。对于整数域逻辑，即使是vshufps，进/出混洗单元也没有转发延迟，但结果应该出现在 FP 域中，为vcmpps 做好准备。使用 vpand 而不是 vandps 对于吞吐量至关重要。

Timothy Furtak 的论文提出了一种使用有效负载对键进行排序的方法：不要将有效负载指针与键打包，而是从比较中生成一个掩码，并以相同的方式在键和有效负载上使用它。这意味着您必须在数据结构中或每次加载结构时将有效负载与键分开。

参见他论文的附录（图 12）。他在键上使用标准的最小值/最大值，然后使用cmpps 查看哪些元素已更改。然后，他在 xor-swap 中间对掩码进行 AND 运算，最终仅将有效负载交换为交换的密钥。

Answer 2

不幸的是，原始 AVX 在其 128 位的一半（即 lanes）中的改组非常有限，因此很难对完整的 256 位寄存器的内容进行排序。但是，AVX2 的 shuffle 操作没有这些限制，因此我们可以以向量化的方式执行一种 4 结构。

我将使用this solution 的想法。为了对数组进行排序，我们必须进行足够的元素比较以确定我们需要应用的排列。鉴于没有元素是 NaN，检查每一对不同的元素 a 和 b 是否 a 以及是否 a > b。有了这些信息，我们就可以充分比较任意两个元素，这必须足以确定最终的排序顺序。这是 6 对 32 位元素和两种比较模式，所以我们最终可以在 AVX 中进行两次随机播放和两次比较。如果您绝对确定所有元素都是不同的，那么您可以避免 a > b 比较并减小 LUT 的大小。

为了在寄存器中重新打包元素，我们可以使用_mm256_permutevar8x32_ps。一条指令允许在 32 位粒度上进行任意洗牌。请注意，在代码中，我假设排序键 f 是您的结构的第一个成员（就像@PeterCordes 建议的那样），但是如果您相应地更改改组掩码，您可以轻松地为当前结构使用此解决方案。

执行比较后，我们有两个 AVX 寄存器，其中包含布尔结果作为 32 位掩码。每个寄存器中的前六个掩码很重要，后两个不重要。然后我们想将这些掩码转换为通用寄存器中的一个小整数，用作查找表中的索引。在一般情况下，我们可能必须为其创建完美的散列，但这里没有必要。我们可以使用_mm256_movemask_ps 从 AVX 寄存器中获取通用寄存器中的 8 位整数掩码。由于每个寄存器的最后两个掩码并不重要，我们可以确保它们始终为零。然后生成的索引将在 [0..2^12) 范围内。

最后，我们从具有 4096 个元素的预计算 LUT 中加载一个混洗掩码，并将其传递给 _mm256_permutevar8x32_ps。结果，我们获得了一个 AVX 寄存器，其中包含 4 个正确排序的类型结构。预计算 LUT 是您的家庭作业 =)

这是最终代码：

__m256i lut[4096];    //LUT of 128Kb size must be precomputed
__m256 Sort4(__m256 val) {
    __m256 aaabbcaa = _mm256_permutevar8x32_ps(val, _mm256_setr_epi32(0, 0, 0, 2, 2, 4, 0, 0));
    __m256 bcdcddaa = _mm256_permutevar8x32_ps(val, _mm256_setr_epi32(2, 4, 6, 4, 6, 6, 0, 0));
    __m256 cmpLt = _mm256_cmp_ps(aaabbcaa, bcdcddaa, _CMP_LT_OQ);
    __m256 cmpGt = _mm256_cmp_ps(aaabbcaa, bcdcddaa, _CMP_GT_OQ);
    int idxLt = _mm256_movemask_ps(cmpLt);
    int idxGt = _mm256_movemask_ps(cmpGt);
    __m256i shuf = lut[idxGt * 64 + idxLt];
    __m256 res = _mm256_permutevar8x32_ps(val, shuf);
    return res;
}

Here 可以看到生成的程序集。总共有 14 条指令，其中 2 条用于加载常量 shuffle 掩码，其中 1 条是由于 movemask 结果的无用 32-bit->64-bit 转换。因此，在一个紧密的循环中，它将是 11-12 条指令。 IACA 表示循环中的四个调用在 Haswell 上具有 16.40 个周期的吞吐量，因此它似乎实现了每个调用 4.1 个周期的吞吐量。

当然，128 Kb 的查找表太大了，除非您打算在一批中处理更多的输入数据。可以通过添加完美的散列来减小 LUT 的大小（当然会牺牲速度）。很难说四个元素可以有多少排序，但显然小于 4！ * 2^3 = 192。我认为 256 元素 LUT 是可能的，甚至可能是 128 元素 LUT。通过完美的散列，将两个 AVX 寄存器通过移位和异或组合成一个，然后执行一次 _mm256_movemask_epi8（而不是执行两个 _mm256_movemask_ps 并随后将它们组合起来）可能会更快。