使用 SSE/AVX/AVX2 检查 __m128i 的所有字节是否匹配单个字节

Question

我正在寻找计算以下函数的有效方法：

输入：__m128i data, uint8_t in;

输出：布尔值，指示是否有任何字节data是in。

我基本上是使用它们来为容量为 8 的字节实现空间/时间高效的堆栈。我最有效的解决方案是首先计算__m128i tmp所有字节为 的 a in。然后检查是否有任何字节tmp\xor data是零字节。

Answer 1

是的，AVX2 具有高效的字节广播。pshufb带有全零掩码的SSSE3同样便宜，但您必须创建随机播放控制向量。AVX512BW/F 甚至有单指令vpbroadcastb/w/d/q x/y/zmm, r32。（使用可选的掩码，因此您可以根据需要将部分归零或与现有向量合并，例如，使用一位掩码在某个位置插入。）

幸运的是，编译器在实现时知道如何执行此操作，_mm_set1_epi8因此我们可以将其留给编译器。

然后它只是归结为通常的pcmpeqb/pmovmskb以获得一个整数，该整数将有1一些用于匹配元素的位，您可以对其进行分支。

// 0 for not found, non-zero for found.  (Bit position tells you where).
unsigned contains(__m128i data, uint8_t needle) {
    __m128i k = _mm_set1_epi8(needle);
    __m128i cmp = _mm_cmpeq_epi8(data, k);  // vector mask
    return _mm_movemask_epi8(cmp);          // integer bitmask 
}

正如您所期望的，所有编译器都使用这个 asm ( Godbolt )

contains(long long __vector(2), unsigned char):
    vmovd   xmm1, edi
    vpbroadcastb    xmm1, xmm1
    vpcmpeqb        xmm0, xmm0, xmm1
    vpmovmskb       eax, xmm0
    ret

除了 MSVC，它首先浪费了一条指令movsx eax, dl。（Windows x64 在 RDX 中传递第二个 arg，而 x86-64 System V在 RDI 中传递第一个整数arg。）

没有 AVX2，你会在 SSSE3 或更高版本中得到类似的东西

# gcc8.3 -O3 -march=nehalem
contains(long long __vector(2), unsigned char):
    movd    xmm1, edi

    pxor    xmm2, xmm2
    pshufb  xmm1, xmm2         # _mm_shuffle_epi8(needle, _mm_setzero_si128())

    pcmpeqb xmm0, xmm1
    pmovmskb        eax, xmm0
    ret

或者只使用 SSE2（x86-64 的基线）：

contains(long long __vector(2), unsigned char):
    mov     DWORD PTR [rsp-12], edi
    movd    xmm1, DWORD PTR [rsp-12]    # gcc's tune=generic strategy is still store/reload  /facepalm
    punpcklbw       xmm1, xmm1          # duplicate to low 2 bytes
    punpcklwd       xmm1, xmm1          # duplciate to low 4 bytes
    pshufd  xmm1, xmm1, 0               # broadcast

    pcmpeqb xmm1, xmm0
    pmovmskb        eax, xmm1
    ret

有关的：

• 如何使用 SIMD 比较两个向量并获得单个布尔结果？和许多重复

• 如何使用 SIMD 计算数组中字节的出现次数？

• SIMD / SSE：如何检查所有向量元素是非零（pxor+ ptest+ jcc= 4的uop与pcmpeqb+ pmovmskb+宏稠合test/jcc。= 3周的uop）

• SSE/AVX 寄存器的非零字节索引（查找匹配位置）

• 如何使用 SIMD 计算字符出现次数（类似于 memchr，但使用 AVX2 计算匹配而不是查找第一个。通过有效的计数累积和有效的水平总和。）