有没有一种有效的方法来使用 SIMD 内在函数来获取 SIMD 寄存器中的第一个非零元素？

Posted 2023-02-19

【中文标题】有没有一种有效的方法来使用 SIMD 内在函数来获取 SIMD 寄存器中的第一个非零元素？

【英文标题】：Is there an efficient way to get the first non-zero element in an SIMD register using SIMD intrinsics?

【发布时间】：2017-02-23 06:43:48

【问题描述】：

如标题所示，如果一个 256 位 SIMD 寄存器是：

0 | 0 | 1 | 0 | 0 | 1 | 0 | 0 |

如何有效地获取第一个非零元素的索引（即第一个1 的索引2）？最直接的方法是存入内存，一一检查，但可能成本很高。有什么好主意吗？

【参考方案1】：

PCMPEQB/W/D/Q 对一个全零寄存器获取一个向量，该向量的元素对于零元素全为 1，对于零元素全为零。 PMOVMSKB 将全为或全零的向量转换为整数位掩码。 （或者movmskps 或pd 获得每个 dword 或 qword 1 位，而不是每个字节，如果这使您的位扫描 - > 索引计算更有效，例如如果您想要一个元素偏移量而不是字节偏移量。 ) 反转它（C ~ 运算符，asm NOT 指令）以在位图中为非零元素获取 1 TZCNT 或 BSF 找到第一个（最低）设置位的整数。如果 BSF 的输入全为零，请注意 BSF 的行为。

---

如果只有一个可能的非零值（如 1），则 PCMPEQB 会针对该向量的一个向量，因此您以后无需反转它。

如果是这种情况，请首先考虑将数据存储在位图中，以将缓存占用空间减少 8 倍。然后您只需 TZCNT 阵列的 64 位块。 （或者使用 SIMD 搜索第一个非零向量，然后 TZCNT 搜索它的第一个非零元素，如果您希望在第一个设置位之前有多个零 qword。就像memcmp 用于查找不匹配字节位置。）

---

刚刚注意到内在标签。 asm指令参考手册在每个条目的底部列出了相关的C内在函数，您可以通过asm助记符搜索Intel's intrinsics finder。 （有关链接，请参阅x86 标签 wiki）。

【讨论】：

谢谢@Peter。我认为您的意思是 LZCNT 而不是 TZCNT。 Acutally asm 指令更好，无论如何都要感谢内在信息。正如您所提到的，只有一个可能的非零值，但是您知道如何在装配级别实现cache footprint 问题吗？

@MarZzz：元素 0 的高位（第一个 arg 到 _mm_set_epi8，最后一个 arg 到 _mm_setr_epi8）进入整数掩码的 LSB。 TZCNT / BSF 首先查看低位，因此使用它们从低地址扫描到高地址（如果向量是从内存中加载的）。如果您想从另一个方向扫描，请使用 LZCNT 或 BSR（它们会给出不同的结果）。

@MarZzz：在 asm 中实现位图有什么不明显的地方？对于这个用例，tzcnt rax, [my_bitmap + rsi] 或其他任何东西，看看从 8*rsi 开始的 64 位中是否有任何命中（因为内存仍然是字节寻址的，除非你使用 BT/BTR/BTS 指令，但不要t 因为它们对内存操作数非常慢，请参阅agner.org/optimize)

感谢您解决 TZCNT 问题，但我对缓存问题感到困惑。您的意思是先将256位数据存储到位图中，没有PCMPEQ或PMOVMSKB，然后每64位TZCNT（即执行4条TZCNT指令）位图？如果是这样，TZCNT 会执行 4 次，这样会更快吗？为什么cache footprint 减少了 8 倍？

@MarZzz：不，我的意思是不要让每个字节都是 0 或 1 的向量，而是提前将它们打包成位。如果您不需要扩展格式的数据用于其他用途，请首先将其存储在打包位图中。我假设您有大量元素，一次操作一个向量，在这种情况下，它的缓存占用量是等效位图的 8 倍。