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

Question

我有一个包含多个数据的64位结构体，其中之一是浮点数值：

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

假设我有四个这样的结构体，放在一个 std::array<MyStruct, 4> 中。

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

- user997112

---

这不是已经是这种情况了吗？ - Kerrek SB

@KerrekSB：不，即使icc输出也只是使用标量运算((u)comiss根据比较低浮点元素设置标志)。是的，它使用指令的v版本，因为这是标量FP数学与-march=native编译时的编译方式，以避免需要在其与其他代码之间使用vzeroupper。 - Peter Cordes

@PeterCordes：谢谢！ - Kerrek SB

我更新了我的答案，加入了一些半聪明的想法，但实际上可能不会带来太多加速。 - Peter Cordes

Answer 1

抱歉，这个答案有些凌乱；它不是一次性写成的，而且我有点懒。其中有一些重复。

我有4个独立的想法：

1. 普通排序，但将结构体作为64位单元移动

2. 矢量化的插入排序作为qsort的构建块

3. 排序网络，使用cmpps/blendvpd实现比使用minps/maxps更好。然而，额外的开销可能会降低速度提升。

4. 排序网络：加载一些结构体，然后洗牌/混合以获得一些仅为浮点数的寄存器和一些仅为有效载荷的寄存器。使用Timothy Furtak的技术进行正常的minps/maxps比较器，然后cmpeqps min,orig->在有效载荷上执行掩码异或交换。这样每个比较器可以排序两倍的数据，但需要在比较器之间匹配两个寄存器的洗牌。完成后需要重新插入（但如果将比较器排列成使得内部解包将最终数据放置在正确顺序中的方式，则这很容易通过unpcklps / unpckhps进行）。

这样做还可以避免一些CPU在处理表示非规范化数、NaN或无穷大的有效负载位模式时可能出现的减速，而不需要通过设置MXCSR中的非规范化数为零位来实现。

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

常规排序

使用常规排序算法时，使用64位加载/存储移动整个结构体，并对FP元素进行标量FP比较可以获得至少小幅加速。为了使这个想法尽可能地发挥作用，请先按浮点值对结构体进行排序，然后可以使用movq将整个结构体移到xmm寄存器中，浮点值将位于ucomiss的低32位。然后你（或者一个聪明的编译器）可以使用movq存储结构体。

查看Kerrek SB链接的汇编输出，编译器似乎在高效地复制结构体方面做得不太好：

icc似乎分别移动两个uint值，而不是通过64位加载来捆绑整个结构体。也许它没有压缩结构体？gcc 5.1大部分时间似乎没有这个问题。

加速插入排序

大型排序通常使用分治法，对于足够小的问题则使用插入排序。 插入排序通过复制数组元素一个一个地向前移动，只有当我们找到当前元素所属位置时才停止。因此，我们需要将一个元素与一系列打包的元素进行比较，如果比较为真，则停止。你闻到向量了吗？我闻到了向量。

# 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，可以节省冗余存储。通过移位/重排寄存器中的比较，而不是直接进行标量加载/比较，可能是一种胜利的方法。这可能会导致太多不可预测的分支，我没有看到一种好的方式来使高4位打包到一个寄存器中供vmovups使用，另一个寄存器中供vmovsd使用。

我可能发明了一种插入排序算法，它兼具两种最糟糕的特点：对于小数组而言速度慢，因为在跳出搜索循环后还需要更多的工作；对于大数组而言速度也慢，因为时间复杂度为O(n^2)。然而，如果搜索循环外部的代码能够被改进得更好，那么这个算法可以作为qsort/mergesort的小数组终点非常有用。

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

更新: Timothy Furtak的论文 描述了一种SSE实现方法来对短数组进行排序(用作构建更大排序，例如这个插入排序)。他建议使用SSE产生一个部分有序的结果，然后使用标量操作进行清理。(在大部分已经排序的数组上进行插入排序是快速的)

这就引导我们到:

排序网络

这里可能没有任何加速。Xiaochen、Rocki和Suda仅报告了在Xeon Phi卡上，对于32位（int）元素的单线程归并排序，从标量到AVX-512的速度提升为3.7倍。对于更宽的元素，向量寄存器中可以容纳的元素较少。（对我们来说是4的因素：在256b中有64b元素，而在512b中有32b元素。）他们还利用AVX512掩码只比较一些通道，这是AVX中没有的功能。此外，由于竞争洗牌/混合单元的较慢的比较器函数，我们已经处于更糟糕的状态。

排序网络可以使用SSE / AVX打包比较指令构建。（更常见的是，使用一对min / max指令，有效地进行一组打包的2元素排序。）更大的排序可以通过进行成对排序的操作来构建。 这篇来自东京大学的Tian Xiaochen，Kamil Rocki和Reiji Suda的论文提供了一些真实的AVX代码进行排序（没有有效载荷），并讨论了向量寄存器的复杂性，因为您无法比较在同一个寄存器中的两个元素（因此，必须设计排序网络不需要这样做）。 它们使用pshufd对齐元素以进行下一次比较，以从仅对几个寄存器的数据进行排序构建出更大的排序。

现在，诀窍在于对64位元素进行成对打包，并基于仅半个元素的比较进行排序（即保持有效负载与排序键）。我们可以通过对（键、有效负载）对数组进行排序来潜在地以这种方式对其他内容进行排序，其中有效负载可以是索引或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位不会使你的FP操作变慢，如果它们恰好代表一个denormal或NaN浮点数。我不确定这是否仅在乘法/除法时发生，或者是否会影响比较。

• Intel Haswell：延迟时间: 等待ymm2准备好需要5个时钟周期，等待ymm3需要7个时钟周期。吞吐量为每4个时钟周期执行一次。(p5瓶颈)。
• Intel Sandybridge/Ivybridge：延迟时间: 等待ymm2准备好需要5个时钟周期，等待ymm3需要6个时钟周期。吞吐量为每2个时钟周期执行一次。(p0/p5瓶颈)。
• AMD Bulldozer/Piledriver: (vblendvpd ymm: 2c 延迟时间，2c 循环吞吐量): 延迟时间: 等待ymm2需要4个时钟周期，等待ymm3需要6个时钟周期，或因cmpps和blend之间的旁路延迟而更差。吞吐量为每4个时钟周期执行一次。(向量P1瓶颈)
• AMD Steamroller: (vblendvpd ymm: 2c 延迟时间，1c 循环吞吐量): 延迟时间: 等待ymm2需要4个时钟周期，等待ymm3需要5个时钟周期，或因旁路延迟而更高。吞吐量为每3个时钟周期执行一次 (cmp和blend的向量端口P0/1瓶颈)。

VBLENDVPD 是2个微操作。由于它有3个寄存器输入，因此不能是1个微操作:/。这两个微操作只能在混洗端口上运行。在Haswell上，这仅限于port5。在SnB上，这是p0/p5。(我不知道为什么Haswell与SnB/IvB相比将混洗/混合吞吐量减半了。)

如果AMD设计具有256位宽的矢量单元，则其低延迟FP比较和解码3输入指令的单个宏操作将使其领先。

通常的minps/maxps对的延迟为3和4个周期(ymm2/3)，吞吐量为每2个周期1个(Intel)。(FP加/减/比较单元上的p1瓶颈)。最公平的比较可能是对64位双精度数进行排序。额外的延迟，如果没有多个独立寄存器对可以比较，可能会有所损失。在Haswell上减半的吞吐量将严重影响任何加速。

请注意，在比较操作之间需要进行混洗以使正确的元素排成一行。最小/最大ps保持了混洗端口未使用，但我的cmpps/blendv版本会使它们饱和，这意味着混洗不能与比较重叠，除非作为填补数据依赖关系留下的空隙的一种方法。

通过超线程，可以使用另一个线程来保持其他端口忙碌（例如0/1端口fp乘法/加法单元或整数代码），从而与这个混合瓶颈版本共享核心。

我尝试了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个uops，相比于blendv版本的5个。在最后6个and / andn / or指令中有很多并行性。cmpps具有3个周期的延迟。我认为ymm2将在6个周期内准备好，而ymm3将在7个周期内准备好。(并且可以与对ymm2的操作重叠)。比较器操作后面的指令可能会是洗牌指令，以将数据放置在下一个比较的正确元素中。对于整数域逻辑运算，即使对于vshufps，也没有前向延迟到/从洗牌单元，但结果应该在FP域中输出，准备进行vcmpps。使用vpand而不是vandps对于吞吐量至关重要。

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

请查看他论文的附录（图12）。他使用标准的最小/最大键，并使用cmpps查看哪些元素发生了更改。然后，他会在xor-swap的中间ANDs该掩码，以便仅交换已交换键的有效载荷。

- Peter Cordes

---

你是否曾遇到过“网络排序”算法？我不确定这是否会对你的提案有所帮助。 - user997112

我听说过排序网络，但从未深入研究过。在没有标准库无法解决的排序需求之前，我从未需要进行排序。除非您谈论运行在非共享内存集群上的排序算法。在这种情况下，我阅读了Andrew Tridgell的一些博士论文。https://www.samba.org/~tridge/phd_thesis.pdf - Peter Cordes

完全没有，我只是想知道它是否有助于AVX提案。 - user997112

我查看了http://www.google.ca/search?&q=sse+sorting+network。使用SSE创建排序网络的方法涉及到像`minps`和`maxps`这样的指令，它们在一对寄存器之间进行操作。但是，无法携带与之相关联的数据，并且根据比较不同元素的结果移动数据。 - Peter Cordes

我最终阅读了一些有关排序网络的内容，并对我的回答进行了大幅更新。 - Peter Cordes

Answer 2

很遗憾，原始的AVX在其128位半部分（即lanes）上有非常有限的洗牌能力，因此很难对完整的256位寄存器内容进行排序。然而，AVX2具有无此限制的洗牌操作，因此我们可以以向量化的方式对4个结构体进行排序。

我将使用这个解决方案的想法。为了对数组进行排序，我们必须进行足够的元素比较，以确保确定需要应用的置换。鉴于没有元素是NaN，则检查每对不同元素a和b是否a < b和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;
}

这里可以看到生成的汇编代码。总共有14条指令，其中2条是用于加载常量洗牌掩码，另外一条是由于movemask结果无用的32位->64位转换而产生的。因此，在紧密循环中，将只有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，然后将它们组合在一起）可能更快。

- stgatilov

---

我认为你不需要进行两次比较。如果你不需要排序是稳定的，那么交换两个相等的元素是可以的。你还可以通过使用pmovzxbd来压缩LUT四倍（使用内部函数将该指令与内存操作数一起发出很难，但即使是movq / pmovzx ymm，也可以）。由于索引很小，你甚至可以将洗牌掩码打包成每个元素4位，并在运行时解包（没有pextr将需要多个指令）。 - Peter Cordes