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

Question

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

输入：__m128i data, uint8_t in;

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

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

Answer 1

是的，AVX2 具有高效的字节广播。 带有全零掩码的 SSSE3pshufb同样便宜，但您必须创建随机控制向量。 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 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 与pcmpeqb+pmovmskb+ 宏融合test/jcc= 3 uops。

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

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