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

Sorting 64-bit structs using AVX?

我有一个 64 位结构体,它代表几条数据,其中一条是浮点值:

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

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

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

抱歉,这个答案很乱;它并没有一次全部写完,我很懒。有一些重复。

我有 4 个不同的想法:

  1. 正常排序,但将结构作为 64 位单元移动
  2. 作为 qsort 构建块的矢量化插入排序
  3. 排序网络,比较器实现使用 cmpps / blendvpd 而不是 minps/maxps。不过,额外的开销可能会扼杀加速。

  4. 排序网络:加载一些结构,然后 shuffle/blend 获取一些仅 float 的寄存器和一些仅 payload 的寄存器。使用 Timothy Furtak 的技术进行正常的 minps/maxps 比较器,然后 cmpeqps min,orig -> 对有效载荷进行屏蔽异或交换。这对每个比较器排序两倍的数据,但确实需要在比较器之间的两个寄存器上进行匹配洗牌。完成后还需要重新交错(但是使用 unpcklps / unpckhps 很容易,如果你安排你的比较器以便那些通道内解包将以正确的顺序放置最终数据)。

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

    Furtak 的论文建议在使用向量对大部分内容进行排序后进行标量清理,这将大大减少改组的数量。

正常排序

使用普通排序算法时,通过将整个结构移动到 64 位 loads/stores 并在 FP 元素上进行标量 FP 比较,至少可以获得小幅加速。为了让这个想法尽可能地发挥作用,首先用浮点值排序你的结构,然后你可以 movq 将整个结构放入 xmm reg,然后浮点值会在 ucomiss 的 low32 中。然后你(或者可能是一个聪明的编译器)可以用 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 直到下一次迭代(或突破之后)将保存冗余存储。通过移位/洗牌来在寄存器中进行比较,而不是字面地做标量 load/compares 可能是一个胜利。这可能会以太多不可预测的分支结束,而且我没有看到一个很好的方法来结束高 4 打包在一个 reg 中用于 vmovups,而低 4 打包在另一个 reg 中用于 vmovsd.

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

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

更新:Timothy Furtak's paper 描述了对短数组进行排序的 SSE 实现(用作更大排序的构建块,如此插入排序)。他建议使用 SSE 生成部分排序的结果,然后使用标量操作进行清理。 (对大部分排序的数组进行插入排序很快。)

这导致我们:

排序网络

这里可能没有任何加速。 Xiaochen、Rocki 和 Suda 只报告了 32 位(整数)元素的标量 -> AVX-512 的 3.7 倍加速,对于 Xeon Phi 卡上的单线程归并排序。对于更宽的元素,更少的元素适合向量 reg。 (这对我们来说是 4 倍:256b 中有 64b 个元素,而 512b 中有 32b 个元素。)他们还利用 AVX512 掩码仅比较某些通道,这是 AVX 中不可用的功能。另外,较慢的比较器函数竞争 shuffle/blend 单元,我们的情况已经很糟糕了。

Sorting networks can be constructed using SSE/AVX packed-compare instructions. (More usually, with a pair of min/max instructions that effectively do a set of packed 2-element sorts.) Larger sorts can be built up out of an operation that does pairwise sorts. 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 位不会减慢您的 FP 操作(如果它们恰好表示非正规或 NaN 浮点数)。我不确定这是否只发生在 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 和混合之间存在旁路延迟。 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 和混合)。

VBLENDVPD 是 2 微指令。 (它有 3 个 reg 输入,所以它不能是 1 uop :/)。两个 uops 只能 运行 在 shuffle 端口上。在 Haswell 上,那只是 port5。在 SnB 上,那是 p0/p5。 (IDK 为什么 Haswell 与 SnB/IvB 相比将洗牌/混合吞吐量减半。)

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

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

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

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

我尝试了 Haswell 的另一个版本,它使用按位 AND/OR 操作来混合 "manually"。不过,它最终变慢了,因为在合并之前,两个来源都必须双向屏蔽。

# 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 and/andn/or 条指令中有很多并行性。但是,cmpps 有 3 个周期的延迟。我认为 ymm2 将在 6 个周期内准备就绪,而 ymm3 将在 7 个周期内准备就绪。(并且可以与 ymm2 上的操作重叠)。比较器操作之后的 insn 可能会被洗牌,以将数据放入正确的元素中以供下一次比较。没有转发延迟 to/from 整数域逻辑的洗牌单元,即使是 vshufps,但结果应该在 FP 域中出现,为 vcmpps 做好准备。使用 vpand 而不是 vandps 对吞吐量至关重要。

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

参见他论文的附录(图12)。他在键上使用标准 min/max,然后使用 cmpps 来查看更改了哪些元素。然后他在异或交换中间对掩码进行 AND 运算,最终只交换交换密钥的有效负载。

不幸的是,原始 AVX 在其 128 位的一半(即 通道 )中的改组非常有限,因此很难对完整的 256 位寄存器的内容进行排序。然而,AVX2 的混洗操作没有这样的限制,所以我们可以用向量化的方式执行一种 4 结构。

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

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

执行比较后,我们有两个 AVX 寄存器包含布尔结果作为 32 位掩码。每个寄存器中的前六个掩码很重要,后两个不重要。然后我们想将这些掩码转换为通用寄存器中的一个小整数,用作查找中的索引 table。在一般情况下,我们可能必须为它创建完美的哈希,但这里没有必要。我们可以使用 _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条是加载constant shuffling masks,1条是由于movemask结果无用的32->64-bit转换。所以在一个紧密的循环中,它将是 11-12 条指令。 IACA 表示循环中的四个调用在 Haswell 上具有 16.40 个周期的吞吐量,因此它似乎实现了每次调用 4.1 个周期的吞吐量。

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