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

Question

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

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

如何有效地获取第一个非零元素的索引（即第一个

2

的索引

1

）？最直接的方法就是存入内存并逐一检查，但成本可能会很高。有什么可爱的想法吗？

Answer 1

• PCMPEQB/W/D/Q 针对全零寄存器，得到一个向量，其中零元素全为 1，零元素全为零。

• PMOVMSKB 将全一或全零的向量转换为整数位掩码。 （或者

  movmskps

  或

  pd

  为每个 dword 或 qword 获取 1 位，而不是每个字节，如果这使您的位扫描 -> 索引计算更高效，就像您想要元素偏移量而不是字节偏移量一样。 )

• 反转（C

  ~

  运算符，asm NOT指令）以在位图中为非零元素获取1

• TZCNT 或 BSF 找到第一个（最低）设置位的整数。如果 BSF 的输入全为零，请注意 BSF 的行为。但幸运的是，当输入是 int

  ~bitmask

  时，这不是问题 - 高 16 个零位变成了 1。 （带有

  vpmovmskb ymm

  的 AVX2 版本，用可能的 1 位填充整个

  uint32_t

  可以使用

  ~(uint64_t)bitmask

  ，或者只使用

  tzcnt

  ，因为 AVX2 CPU 也有 BMI1。）

  在 C++20 中，这是

  std::countr_one(mask)

  实现为

  std::countr_zero(~mask)

  。

  std::countr_zero

  与 x86

  tzcnt

  具有相同的语义，包括 input=0。如果您在没有 BMI1 的情况下进行编译，请检查 asm 以确保编译器优化了对 input==0 的任何检查，因为

  uint32_t

  ~mask

  对于 16 位向量始终非零。

---

例如使用内在函数：

int first_nonzero_byte(__m128i v){
    //__m128i v = _mm_loadu_si128((const __m128i*)p);  // for a pointer arg
    __m128i vcmp = _mm_cmpeq_epi8(v, _mm_setzero_si128());
    unsigned bitmask = _mm_movemask_epi8(vcmp);
#ifdef __GNUC__
    return __builtin_ctz(~bitmask);
#else
    return _tzcnt_u32( ~bitmask );
#endif
   // returns 16 if v was all zero so ~bitmask is 0xFFFF0000
}

在 https://godbolt.org/z/Y8vYbsW69 上编译至

# GCC11.2 -O3 -msse4.1
        movdqa  xmm1, xmm0      # missed optimization, should zero XMM1 instead
        pxor    xmm0, xmm0
        pcmpeqb xmm0, xmm1
        pmovmskb        eax, xmm0
        not     eax
        rep bsf eax, eax        # tzcnt on new CPUs, BSF on old
        ret

在 GNU C 中，如果没有

_tzcnt_u32

或其他东西，

-march=haswell

将无法编译，我们使用

__builtin_ctz

。正如我所说，

~bitmask

保证非零。

tzcnt

编码为

rep bsf

；旧的 CPU 会将其执行为

bsf

，对于非零输入产生相同的结果。新的 CPU 会将其执行为

tzcnt

，这在 AMD 上效率更高（2 uops 而不是 7）。英特尔以单微指令的形式执行其中之一。如果您不告诉 GCC 要调整的特定 CPU，GCC 将使用

rep bsf

又名

tzcnt

。

对于 JATothrim 的答案中所示的相关功能，仅使用 4 个单微指令（实际上 AMD 上的 tzcnt 为 2 个微指令），而不是 8 个指令，包括

pblendvb

（英特尔上的 2 个微指令）。如果您希望元素索引作为

vpermilps

的洗牌控制向量，则该答案中的洗牌/水平缩减想法可能很有用，但当您实际上想要标量

int

时，与此相比似乎不是最佳选择。

int equal_first_dword_bitscan(__m128i x, __m128i y)
{
    __m128i vcmp = _mm_cmpeq_epi32(x,y);
    unsigned bitmask = _mm_movemask_ps(_mm_castsi128_ps(vcmp));
    bitmask |= 1<<4;    // return 4 if the low 4 bits are all 0
#ifdef __GNUC__
    return __builtin_ctz(bitmask);
#else
    return  _tzcnt_u32( bitmask );  // runs as BSF on old CPUs, don't skip the OR
#endif
}

MSVC 没有

__builtin_ctz

，但即使你没有告诉它目标 CPU 支持 BMI1，它也会编译

_tzcnt_u32

。如果您确实只在 BMI1 的 CPU 上运行，则可以省略

bitmask |= 1<<4;

，这样它将返回

32

表示未找到。

如果您在多个函数中使用尾随零计数，最好将 ifdef 内容包装在辅助函数中，而不是在每个用例中。

---

如果只有一个可能的非零值（如

1

），则 PCMPEQB 针对该向量，这样您以后就不需要反转它。

如果是这种情况，请考虑首先将数据存储在位图中，以将缓存占用空间减少 8 倍。然后只需 TZCNT 64 位数组块即可。

或者对于更大的数据数组，使用 SIMD 搜索第一个非零向量，然后 TZCNT 搜索其中的第一个非零元素（如果您希望在第一个设置位之前有多个 qwords 的零）。就像

memcmp

那样查找不匹配的字节位置。
请参阅高效查找大型数组中的最低有效设置位？和如何高效查找数组中的第一个非零？

---

顺便说一句，asm 指令参考手册在每个条目的底部列出了相关的 C 内在函数，您可以通过 asm 助记符搜索 Intel 的内在函数查找器。 （请参阅 x86 标签 wiki 获取链接）。

Answer 2

我最近一直在编写一堆“获取 X 的索引”SIMD 算法。 到目前为止，从比较掩码中提取索引的最通用方法是通过水平indice最小值。

这是（无符号）整数水平最小值：

int horizontal_min(__m128i x) {
    x = _mm_min_epu32(x, _mm_shuffle_epi32(x, 0b01001110));
    x = _mm_min_epu32(x, _mm_shuffle_epi32(x, 0b11100001));
    return _mm_extract_epi32(x,0);
}

现在执行以下操作：

int equal_first(__m128i x, __m128i y) {
    const __m128i index = _mm_set_epi32(0,1,2,3);
    // Compute mask
    __m128i mask = _mm_cmpeq_epi32(x,y);
    // Select indices.
    mask = _mm_blendv_epi8(_mm_set1_epi32(-1), index, mask);
    // mask = index | (~mask);
    // pick smallest indice.
    return horizontal_min(mask);
}

这段代码的优点是不需要任何位扫描指令，全部在FPU上完成。

提示：如果您使用

phminposuw128

指令来计算最小索引，那么 16 位索引会变得非常高效。

编辑：Peter 的分析指出我的解决方案速度较慢，除非您需要 SIMD 寄存器中的结果。

另一种情况是缩减循环，您需要数组中所述元素的索引。 在循环中，您累积了例如SIMD 寄存器中的最小/最大元素索引。现在无序的索引可能指向源数组中的任何位置。现在你必须使用horizontal_min()来告诉最小/最大元素在哪里。