使用 simd 在双精度数组中查找 nan

Posted 2023-02-16

【中文标题】使用 simd 在双精度数组中查找 nan

【英文标题】：find nan in array of doubles using simd

【发布时间】：2020-05-24 05:21:14

【问题描述】：

这个问题非常类似于：

SIMD instructions for floating point equality comparison (with NaN == NaN)

虽然该问题侧重于 128 位向量，并且有关于识别 +0 和 -0 的要求。

我有一种感觉，我自己可以得到这个，但英特尔内在函数指南页面似乎已关闭：/

我的目标是获取一个双精度数组并返回数组中是否存在 NaN。我预计大多数时候不会有一个，并且希望该路线具有最佳性能。

最初，我打算对自己的 4 个双精度进行比较，以反映用于 NaN 检测的非 SIMD 方法（即仅当 a != a 为真时的 NaN 值）。比如：

data *double = ...
__m256d a, b;
int temp = 0;

//This bit would be in a loop over the array
//I'd probably put a sentinel in and loop over while !temp
a = _mm256_loadu_pd(data);
b = _mm256_cmp_pd(a, a, _CMP_NEQ_UQ);
temp = temp | _mm256_movemask_pd(b);

但是，在某些比较示例中，除了比较本身之外，似乎还进行了某种 NaN 检测。我简单地想，如果像 _CMP_EQ_UQ 这样的东西会检测到 NaN，我可以使用它，然后我可以将 4 个双打与 4 个双打进行比较，并神奇地同时查看 8 个双打。

__m256d a, b, c;
a = _mm256_loadu_pd(data);
b = _mm256_loadu_pd(data+4);
c = _mm256_cmp_pd(a, b, _CMP_EQ_UQ);

此时我意识到我并没有完全正确地思考，因为我可能碰巧将一个不是 NaN 的数字与其自身进行比较（即 3 == 3）并以这种方式获得成功。

所以我的问题是，将 4 个双精度数与自己进行比较（如上所述）是我能做的最好的，还是有其他更好的方法来找出我的数组是否有 NaN？

【参考方案1】：

您可以通过检查 fenv 状态来完全避免这种情况，或者如果没有，则缓存阻止它和/或将其折叠到相同数据的另一个通道中，因为它的计算强度非常低（加载/存储的每个字节的工作量） ，所以很容易成为内存带宽的瓶颈。见下文。

---

您要查找的比较谓词是 _CMP_UNORD_Q 或 _CMP_ORD_Q，告诉您比较是无序的还是有序的，即至少有一个操作数是 NaN，或者两个操作数分别是非NaN。 What does ordered / unordered comparison mean?

cmppd 的 asm 文档列出了谓词，并且具有与内在函数指南相同或更好的详细信息。

所以是的，如果您希望 NaN 很少见并且想要快速扫描大量非 NaN 值，您可以 vcmppd 两个不同的向量相互对抗。如果您关心 NaN 在哪里，那么一旦您知道两个输入向量中的任何一个中至少有一个，您就可以做额外的工作来解决这个问题。 （就像_mm256_cmp_pd(a,a, _CMP_UNORD_Q) 提供移动掩码 + 位扫描以获得最低设置位。）

---

OR 或 AND 多次比较每个 movemask

与其他 SSE/AVX 搜索循环一样，您还可以通过将一些比较结果与 _mm256_or_pd（查找任何未排序的）或 _mm256_and_pd（检查所有已排序）相结合来分摊 movemask 成本。例如。 每个移动掩码/测试/分支检查几个缓存行（4x _mm256d 和 2x _mm256_cmp_pd）。（glibc 的 asm memchr 和 strlen 使用这个技巧。）同样，这优化了您的常见情况是您不希望提前退出并且必须扫描整个阵列。

还请记住，检查同一个元素两次完全没问题，因此您的清理工作可以很简单：加载到数组末尾的向量，可能与您已经检查过的元素重叠。

// checks 4 vectors = 16 doubles
// non-zero means there was a NaN somewhere in p[0..15]
static inline
int any_nan_block(double *p) 
    __m256d a = _mm256_loadu_pd(p+0);
    __m256d abnan = _mm256_cmp_pd(a, _mm256_loadu_pd(p+ 4), _CMP_UNORD_Q);
    __m256d c = _mm256_loadu_pd(p+8);
    __m256d cdnan = _mm256_cmp_pd(c, _mm256_loadu_pd(p+12), _CMP_UNORD_Q);
    __m256d abcdnan = _mm256_or_pd(abnan, cdnan);
    return _mm256_movemask_pd(abcdnan);

// more aggressive ORing is possible but probably not needed
// especially if you expect any memory bottlenecks.

我把 C 写成汇编，每个源代码行一条指令。 （加载/内存源 cmppd）。如果在 Intel 上使用非索引寻址模式，这 6 条指令在现代 CPU 的融合域中都是单微指令。 test/jnz 作为 break 条件将使其达到 7 微秒。

在循环中，add reg, 16*8 指针增量是另一个 1 uop，而 cmp / jne 作为循环条件是另一个，使其达到 9 uop。因此不幸的是，在 Skylake 上，前端的瓶颈在 4 微秒 / 时钟，至少需要 9/4 个周期才能发出 1 次迭代，而负载端口并没有完全饱和。 Zen 2 或 Ice Lake 可以在每个时钟维持 2 次负载，而无需再展开或其他级别的 vorpd 组合。

---

另一个可能的技巧是在两个向量上使用 vptest 或 vtestpd 来检查它们是否都非零。但我不确定是否可以正确检查两个向量的 每个 元素是否非零。 Can PTEST be used to test if two registers are both zero or some other condition? 表明另一种方式（_CMP_UNORD_Q 输入均为全零）是不可能的。

但这并没有真正的帮助：vtestpd / jcc 总共是 3 个微指令，而 vorpd / vmovmskpd / test+jcc 在现有的 Intel/AMD CPU 上也是 3 个融合域微指令AVX，因此当您在结果上进行分支时，它甚至不是吞吐量的胜利。因此，即使有可能，它也可能是收支平衡的，尽管它可能会节省一些代码大小。如果需要多个分支从全零案例中分出全零或 mix_zeros_and_ones 案例，则不值得考虑。

---

避免工作：改为检查 fenv 标志

如果您的数组是该线程中计算的结果，只需检查 FP 异常粘滞标志（手动在 MXCSR 中，或通过 fenv.hfegetexcept）查看自上次以来是否发生了 FP“无效”异常清除 FP 异常。如果不是，我认为这意味着 FPU 没有产生任何 NaN 输出，因此从那时起该线程写入的数组中没有任何输出。

如果已设置，则必须检查；可能已针对未传播到此数组的临时结果引发了无效异常。

---

缓存阻塞：

如果/当 fenv 标志不能让您完全避免这项工作，或者对您的程序来说不是一个好的策略，尝试将此检查折叠到生成数组的任何内容中，或者放入读取的下一个通道中它。因此，当数据已经加载到向量寄存器中时，您正在重用数据，从而增加了计算强度。 （每次加载/存储的 ALU 工作量。）

即使数据在 L1d 中已经很热，它仍然会成为负载端口带宽的瓶颈：每个 cmppd 2 个负载仍然是 2 个时钟负载端口带宽的瓶颈，在具有 2 个时钟 vcmppd ymm 的 CPU 上（Skylake 但不是哈斯韦尔）。

对齐指针以确保从 L1d 缓存中获得满负载吞吐量也是值得的，尤其是当数据有时在 L1d 中已经很热时。

或者至少 缓存块，这样您就可以在缓存中热时在同一块上运行另一个循环之前检查一个 128kiB 块。这是 256k L2 大小的一半，因此您的数据在上一次传递中应该仍然很热，并且/或者在下一次传递中仍然很热。

绝对避免在整个数兆字节的数组上运行它，并支付将其从 DRAM 或 L3 缓存进入 CPU 内核的成本，然后在另一个循环读取它之前再次驱逐。这是最坏情况下的计算强度，需要付出不止一次将其放入 CPU 内核的私有缓存的成本。

【讨论】：

完美。我不得不通读您的答案几次，以尝试了解所有有用的花絮。我之前没有考虑过重新处理数组的末尾。在类似的代码中，我添加了 openmp 并发现与 SIMD 一起使用时没有任何好处，经过一些数字运算后我意识到这是因为内存带宽限制。

@Jimbo：在“客户端”芯片（台式机/笔记本电脑）上，单核几乎可以使内存带宽饱和。在多核 Xeon 芯片上情况并非如此：它们具有更多的总带宽和更低的单核带宽。 Why is Skylake so much better than Broadwell-E for single-threaded memory throughput?。就像我说的那样，这个检查的计算强度非常低，应该折叠到数据的不同通道中，或者至少缓存阻塞，这样数据在 L2 或 L1d 缓存中是热的，用于这个或下一个通道。