在 4 个 __m256d 寄存器中找到 4 个最小值

Posted 2023-02-16

【中文标题】在 4 个 __m256d 寄存器中找到 4 个最小值

【英文标题】：Find 4 minimal values in 4 __m256d registers

【发布时间】：2016-03-11 16:28:11

【问题描述】：

我不知道如何实现：

__m256d min(__m256d A, __m256d B, __m256d C, __m256d D)

    __m256d result;

    // result should contain 4 minimal values out of 16 : A[0], A[1], A[2], A[3], B[0], ... , D[3]
    // moreover it should be result[0] <= result[1] <= result[2] <= result[2]

     return result;

对如何以智能方式使用_mm256_min_pd、_mm256_max_pd 和随机/置换有任何想法吗？

================================================ ===

这是我到目前为止的地方，之后：

    __m256d T = _mm256_min_pd(A, B);
    __m256d Q = _mm256_max_pd(A, B);
    A = T; B = Q;
    T = _mm256_min_pd(C, D);
    Q = _mm256_max_pd(C, D);
    C = T; D = Q;
    T = _mm256_min_pd(B, C);
    Q = _mm256_max_pd(B, C);
    B = T; C = Q;
    T = _mm256_min_pd(A, B);
    Q = _mm256_max_pd(A, B);
    A = T; D = Q;
    T = _mm256_min_pd(C, D);
    Q = _mm256_max_pd(C, D);
    C = T; D = Q;
    T = _mm256_min_pd(B, C);
    Q = _mm256_max_pd(B, C);
    B = T; C = Q;

我们有： A[0]

所以最小值在 A 之间，第二个最小值在 A 或 B 之间，... 不知道从那里去哪里......

================================================ =========

第二个想法是问题可以简化为自身，但有 2 个输入 __m256 元素。如果可以做到，那么只需执行 min4(A,B) --> P, min4(C,D) --> Q, min4(P,Q) --> 返回值。

不知道如何处理两个向量：)

================================================ =========================

更新 2：问题几乎解决了——以下函数计算 4 个最小值。

__m256d min4(__m256d A, __m256d B, __m256d C, __m256d D)

    __m256d T;
    T = _mm256_min_pd(A, B);
    B = _mm256_max_pd(A, B);            
    B = _mm256_permute_pd(B, 0x5);
    A = _mm256_min_pd(T, B);            
    B = _mm256_max_pd(T, B);            
    B = _mm256_permute2f128_pd(B, B, 0x1);
    T = _mm256_min_pd(A, B);
    B = _mm256_max_pd(A, B);
    B = _mm256_permute_pd(B, 0x5);
    A = _mm256_min_pd(A, B);

    T = _mm256_min_pd(C, D);
    D = _mm256_max_pd(C, D);            
    D = _mm256_permute_pd(D, 0x5);
    C = _mm256_min_pd(T, D);            
    D = _mm256_max_pd(T, D);            
    D = _mm256_permute2f128_pd(D, D, 0x1);
    T = _mm256_min_pd(C, D);
    D = _mm256_max_pd(C, D);
    D = _mm256_permute_pd(D, 0x5);
    C = _mm256_min_pd(C, D);

    T = _mm256_min_pd(A, C);
    C = _mm256_max_pd(A, C);            
    C = _mm256_permute_pd(C, 0x5);
    A = _mm256_min_pd(T, C);            
    C = _mm256_max_pd(T, C);            
    C = _mm256_permute2f128_pd(C, C, 0x1);
    T = _mm256_min_pd(A, C);
    C = _mm256_max_pd(A, C);
    C = _mm256_permute_pd(C, 0x5);
    A = _mm256_min_pd(A, C);

    return A;
;

剩下的就是在返回之前对 A 内部的值进行升序排序。

【问题讨论】：

您遇到的具体问题是什么？这是一个相当广泛的问题。

您正在寻找一个将所有 16 个双精度数中最低的 4 个双精度数按顺序排列到一个向量中，对吗？谷歌 SIMD 排序网络之类的东西。您可能会发现解包到两个 __m128d 向量对于某些步骤很有用，但可能不是。如果您只关心最小的 4 个元素，而不是完整的排序，那么使用 SIMD 排序网络可能更难击败标量代码。

正确——所有 16 个双精度数中最低的 4 个双精度数成为一个向量。这 4 个向量包含 16 个值，这些值是 SIMD 计算的结果，效果非常好。最后必须选择最低的 4 个。目标不是击败标量代码，而只是避免它。将值卸载到内存，然后进行排序，然后再次加载，对我来说似乎不合逻辑。

这里最重要的性能标准是什么？延迟、吞吐量或总融合域微指令（即对周围代码吞吐量的影响）？乱序执行是否可能同时进行多个排序，或者与其他工作重叠，或者这是循环携带依赖的一部分？

顺便说一句，从 n 中找出最小的 k 个元素的正式名称是 Selection algorithm。 （从技术上讲，选择算法只找到第 k 阶统计量，而不是所有 k..n 或 0..k（部分排序）。我们想要一个不做额外工作的部分排序来确保数组的其余部分（或寄存器）仍然有有意义的数据。）无论如何，我没有找到很多关于 very small k 的讨论，其中 n 在谷歌搜索 simd selection algorithm 时也很小。 ://

【参考方案1】：

最好做一些 SIMD 比较以减少到 8 或 4 个（就像你现在有）候选，然后解包到向量寄存器中的标量双精度数。这不必涉及内存往返：vextractf128 高半部分 (_mm256_extractf128_pd)，然后投射低半部分。也许使用movhlps (_mm_movehl_ps) 将__m128 的高半部分降低到低半部分（尽管在带有 AVX 的 CPU 上，您只需要保存一两个代码字节即可使用它，而不是立即使用 shuffle ; 它并不像在某些旧 CPU 上那样快）。

IDK 无论是通过随机播放打开包装还是简单地存储都是可行的方法。也许两者兼而有之，以保持洗牌端口和存储/加载端口繁忙会做得很好。显然，每个向量中的低 double 已经作为标量存在，因此您不必加载它。 （而且编译器不善于通过存储并重新加载为标量来利用这一点，即使是本地数组也是如此。）

即使没有非常缩小候选集，解包前的一些 SIMD 比较器也可以减少分支标量代码预期的交换/洗牌量，从而减少分支错误预测的惩罚。

---

正如我在 cmets 关于 Paul R 的回答中所描述的，在标量代码中，您可能会使用插入排序类型的算法做得很好。但它更像是一个优先队列：只插入前 4 个元素。如果一个新的候选人比最大的现有候选人更大，那就继续吧。否则，将其插入排序到您按排序顺序维护的 4 个候选者列表中。

---

我找到了really nice paper on SIMD sorting networks, with specific discussion of AVX。他们详细介绍了使用 SIMD packed-min / packed-max 指令对几个数据向量寄存器进行排序时所需的洗牌。他们甚至在他们的示例中使用像 _mm512_shuffle_epi32 这样的内在函数。他们说他们的结果适用于 AVX，即使他们在示例中使用了 AVX-512 掩码寄存器。

这只是论文的最后一点，他们讨论了合并以使用小型排序作为大型并行排序的构建块。我在任何地方都找不到他们的实际代码，所以也许他们从未发布过他们为制作图表而进行基准测试的完整实现。 :(

顺便说一句，我之前写过一篇answer with some not-very-great ideas，关于通过float 成员对 64 位结构进行排序，但这在这里并不适用，因为我只是解决了处理有效负载（你没有）的复杂性。

---

我现在没有时间完成这个答案，所以我将发布我的想法的摘要：

将该论文中的 2-register 方法改编为 AVX（或 AVX2）。我不确定如何最好地模拟他们的 AVX512 屏蔽最小/最大指令。 ：/我可能会稍后更新。您可能想给作者发电子邮件，询问他们用于对桌面 CPU 进行基准测试的代码。

无论如何，在成对的 reg 上使用 2-register 功能，将 4 个 reg 减少到 2 个 reg，然后再次减少到 1 个 reg。与您的版本不同，他们的版本生成一个完全排序的输出寄存器。

尽可能避免跨车道洗牌可能会很棘手。我不确定您是否可以通过使用 shufpd (__m256d _mm256_shuffle_pd (__m256d a, __m256d b, const int select);) 在改组时合并来自两个源 reg 的数据来获得任何好处。 256b 版本可以在每个通道上进行不同的 shuffle，使用 imm8 的 4 位而不是 2 位。

这是一个有趣的问题，但不幸的是，我不应该花时间自己编写完整的解决方案。如果我有时间，我想比较一个插入排序优先级队列和一个排序网络完全展开的同一 pqueue 实现，每个有 4、8、12 和 16 个元素。 （标量之前的不同级别的 SIMD 排序网络）。

我找到的链接：

Another paper on SSE sorting，一些代码使用palignr 将两个两个单独排序的向量合并为一个排序的 8 元素向量对。不直接适用于双精度的 256b 向量。 Another paper on SSE sorting，来自 Core2 CPU 的测试结果。由于shufps 的限制，它使用低效的混合/混合/洗牌在两个向量之间进行洗牌。 In-lane shufpd 的限制略有不同。

这篇论文可能值得仔细阅读

。他们拥有适用于真实 SSE 向量的算法，并具有可用的 shuffle 操作。 Xiaochen, Rocki, and Suda's paper linked already What looks like a chapter on sorting networks 来自the CLR algorithms textbook。 Sorting network generator，可选择算法。不是特定于 SIMD 的 https://en.wikipedia.org/wiki/Bitonic_sorter 示例图有一个 16 输入排序网络。双调排序使用在某种程度上可以进行 SIMD 的模式。网络末端的上半部分可以省略，因为我们只关心最低 4 的顺序。 Fastest sort of fixed length 6 int array。一个有很多答案的热门问题。

【讨论】：

好答案！我将在一篇我认为看起来很有希望的论文中添加另一个链接。它在第 6 页 (1279) vldb.org/pvldb/vol8/p1274-inoue.pdf

上有伪代码

@gustf：不仅仅是伪代码：具有内在函数的实际 C++。有趣：我一直忘记palignr 用于组合来自两个向量的元素。当然，这个问题是在问浮点数，所以 palignr 会导致额外的延迟转发到 minpd/maxpd。他们用它来传输单个元素，所以很遗憾，它没有映射到_mm256_permute2f128_pd。

没错，我的意思是写内在函数，不知道到底发生了什么 :) 是的，你对 alignr 的看法是对的，但我认为实际的算法可能会引起人们的兴趣。

【参考方案2】：

这是一个纯粹的“水平”操作，并不真正适合 SIMD - 我怀疑将四个向量存储在内存中，对 16 个值进行排序，然后将前四个加载到结果向量中会更快：

__m256d min(__m256d A, __m256d B, __m256d C, __m256d D)

    double buff[16] __attribute__ ((aligned(32)));

    _mm256_store_pd(&buff[0], A);
    _mm256_store_pd(&buff[4], B);
    _mm256_store_pd(&buff[8], C);
    _mm256_store_pd(&buff[12], D);

    std::partial_sort(buff, buff+4, buff+16);

    return _mm256_load_pd(&buff[0]);    

为了提高性能，您可以实现一个内联自定义排序例程，该例程针对 16 个元素进行硬编码。

【讨论】：

使用std::partial_sort(buff, buff+4, buff+16) 不会浪费时间对整个数组进行排序。

在我查看之前我不确定它是否有一个 STL 函数，但我知道这个概念是存在的。希望当部分排序范围非常小时，它有不同的策略。例如插入排序到前 4 个元素中，然后停止检查。或者维护到目前为止看到的最高 4 个元素的队列。 std::partial_sort 仍然做了比需要更多的工作，因为它不能让数组的其余部分损坏（例如元素的重复副本）。也许有一个 STL 函数更适合这个，但 partial_sort 是我首先发现的。

另外，对于-std=gnu++11，您可以使用alignas(32) double buff[16];。然后 gcc 和 clang 在设置堆栈帧后生成必要的 and rsp, -32 or equivalent。不幸的是，它看起来并不高效：/ 可能是一些矢量排序网络的东西，然后水平标量超过 4 或 8 个元素会更好。

@Fedor_Govnjukoff：插入排序被认为是小型数组的最佳选择。

@Fedor_Govnjukoff：适应插入排序：在前 4 个元素上正常运行，因此它们被排序。之后，只考虑插入前 4 个元素。当移动为新元素腾出空间时，您可以在编写新的第 4 个元素后停止移动。因此，您基本上将前 4 个元素视为您插入的优先级队列。对于后面的每个元素，首先检查 pqueue 中最大的元素，看看该元素是否小于当前 min4 中的任何一个。在具有低延迟 gpvector 的 Intel CPU 上，可能是广播负载和 cmpps->movmsk->bsf