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

时间:2019-02-27 02:54:38

标签: sse simd avx avx2

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

输入:__m128i data, uint8_t in;

输出:布尔值,指示data中的任何字节是否为in

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

1 个答案:

答案 0 :(得分:3)

是的,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中传递了第二个参数,而x86-64系统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

相关: