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

Question

我正在寻找有效的方法来计算以下功能：

输入：__m128i data, uint8_t in;

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

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

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 
}

正如您所料，所有编译器都使用此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

---

有关：

• How to compare two vectors using SIMD and get a single boolean result?和许多重复
• How can I count the occurrence of a byte in array using SIMD?
• SIMD/SSE: How to check that all vector elements are non-zero（pxor + ptest + jcc = 4 uops vs. pcmpeqb + pmovmskb + macro-fused test/jcc = 3 uops。）
• The indices of non-zero bytes of an SSE/AVX register（找到比赛位置）
• How to count character occurrences using SIMD（就像memchr一样，但计算匹配而不是找到第一个，使用AVX2。有效累积计数，以及有效的水平和。）