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

Question

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

输入：__m128i data，uint8_t in；

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

我主要使用它们来实现一个具有8个字节容量的空间/时间高效的字节堆栈。我的最有效解决方案是首先计算出一个__m128i tmp，其中所有字节都是in。然后检查tmp\xor data中是否有任何字节为零字节。

- redplum

Answer 1

是的，AVX2具有高效的字节广播功能。SSSE3中使用全零掩码的pshufb操作同样廉价，但您需要创建洗牌控制向量。甚至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 
}

如您所料，所有的编译器都使用这个汇编代码（在

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 寄存器；而在 x86-64 System V 中，第一个整数参数传递到 RDI 寄存器。)

---

如果没有 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 uops vs. pcmpeqb+pmovmskb +宏融合的test/jcc=3 uops。）

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

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

- Peter Cordes

---

不错！我不知道还有字节广播指令！ - redplum