有没有一种有效的方法可以使用 SIMD 内部函数获取 SIMD 寄存器中的第一个 non-zero 元素？

Question

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

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

如何有效地获取第一个 non-zero 元素的索引（即第一个 1 的索引 2）？最直接的方法就是存入内存，逐一检查，但代价太大。有什么可爱的点子吗？

Answer 1

PCMPEQB/W/D/Q 针对全零寄存器得到一个向量，其元素对于零元素全为 1，对于零元素全为零。
PMOVMSKB 将全 1 或全 0 的向量转换为整数位掩码。（或者 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 CPUs也有 BMI1。)

例如内在函数：

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；旧的 CPUs 将执行它作为 bsf，为非零输入产生相同的结果。新的 CPUs 将执行它作为 tzcnt，这在 AMD 上更有效（2 微指令而不是 7）。英特尔作为单 uop 执行。 GCC 使用 rep bsf aka tzcnt 如果你不告诉它一个特定的 CPU 来调整。

对于 JATohrim 的回答中所示的相关功能，仅使用 4 条单指令（实际上 AMD 上的 tzcnt 为 2 指令）而不是 8 条指令，包括 pblendvb（英特尔为 2 指令）。如果您希望元素索引作为 vpermilps 的混洗控制向量，那么该答案中的 shuffle/horizontal-reduction 想法可能会有用，但当您真正想要标量 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，但会编译 _tzcnt_u32，即使您没有告诉它目标 CPU 支持 BMI1。如果你肯定只有运行在 CPUs 与 BMI1，你可以省略 bitmask |= 1<<4; 这样它会 return 32 未找到。

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

如果只有一个可能的非零值（如 1），则 PCMPEQB 会针对该值的一个向量，这样您以后就不需要反转它了。

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

或者对于更大的数据数组，用 SIMD 搜索第一个非零向量，然后 TZCNT 搜索它的第一个非零元素，如果你期望在那里在第一个设置位之前是多个零的 qwords。就像 memcmp 查找不匹配的字节位置一样。
参见 and How to find the first nonzero in an array efficiently?

顺便说一句，asm 指令参考手册在每个条目的底部列出了相关的 C 内在函数，您可以在 Intel's intrinsics finder by asm mnemonic. (See the x86 标签 wiki 中搜索链接。

Answer 2

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

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

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 寄存器中的结果，否则我的解决方案速度较慢。

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

有没有一种有效的方法可以使用 SIMD 内部函数获取 SIMD 寄存器中的第一个 non-zero 元素？

Is there an efficient way to get the first non-zero element in an SIMD register using SIMD intrinsics?

x86

bit-manipulation

simd

intrinsics

avx