SIMD 中将模式与位掩码进行比较的最快算法是什么？

Question

我想在SIMD中优化以下代码

拍：[1,2,3,4]

数据：[1,1,3,3]

mask: [1, 0, 1, 1] # 1表示相等，0表示可选/不关心

结果：[1, 1, 1, 0]

结果的朴素算法：

  for i in range(len(pat)):
    result[i] = mask[i] ? pat[i] == data[i] : 1

结果的替代朴素算法：

  for i in range(len(pat)):
    result[i] = !mask[i] || pat[i] == data[i]

这是我使用 Rust 便携式 SIMD 所做的事情：

fn simd_match<const N: usize>(pattern: [u8; N], mask: u64, data: [u8; N]) -> Mask<i8, N>
where
    LaneCount<N>: SupportedLaneCount,
{
    Mask::from_bitmask(mask)
        .select_mask(
            Simd::from_array(data).simd_eq(Simd::from_array(pattern)),
            Mask::from_bitmask(u64::MAX),
        )
}

有什么办法可以改善这个问题吗？我想知道在不同CPU情况下执行此操作的最佳方法，例如如果 AVX512BW 可用，则可以使用 vpcmpb/kmovd。

我尝试在 Rust 中使用进行运行时检测的多版本功能来实现它，但我不确定我的目标选择是否足够：

#[inline(always)]
#[multiversion(targets(
    "x86_64+avx512bw",
    "x86_64+avx2",
    "x86_64+avx",
    "x86_64+sse4.2",
    "x86_64+sse2",
    "x86+avx512bw",
    "x86+avx2",
    "x86+avx",
    "x86+sse4.2",
    "x86+sse2",
    "aarch64+sve2",
    "aarch64+sve",
    "aarch64+neon",
    "arm+neon",
    "arm+vfp4",
    "arm+vfp3",
    "arm+vfp2",
    
))]

fn simd_match<const N: usize>(pattern: [u8; N], mask: u64, data: [u8; N]) -> bool
where
    LaneCount<N>: SupportedLaneCount,
{
    Mask::from_bitmask(mask)
        .select_mask(
            Simd::from_array(data).simd_eq(Simd::from_array(pattern)),
            Mask::from_bitmask(u64::MAX),
        )
        .all()
}

Answer 1

对于没有 AVX-512 的 x86，您想要的 asm 是

[v]pcmpeqb

/

[v]pmovmskb

以获得标量掩码，然后是标量整数

and

（可能在将两到四个 16 位掩码组合成 32 位或 64 位之后）带移位/或的位寄存器。

您的内在函数看起来像是将

u64

掩码转换为向量，这明显不太方便，特别是如果您最终想要或可以使用位掩码而不是向量。 （intel avx2 中是否有与 movemask 指令相反的指令？ - 不，你必须用一些指令来模拟它）。

使用 AVX-512 零掩码比较掩码是一个免费执行 AND 操作的选项，如

vpcmpeqb k1{k1}, zmm0, [rdi]

（

z

上的正常

{}

后缀是隐式的；比较仅支持零掩码而不支持合并） -masking，但不幸的是，asm 语法需要省略

z

）。 如果您想将其与较窄的向量一起使用，则内部函数

_mm512_mask_cmp_epi8_mask

或

_mm256_...

或 128。

屏蔽将抑制内存源操作数加载时出现的错误；对于内在函数，您可能需要使用与比较具有相同掩码的掩码加载，但希望普通的

_mm512_loadu_si512

仍然可以折叠到内存源操作数中进行掩码比较，因为越界访问是 UB，故障不是所需的副作用。

但对于 32 字节或更小的向量，即使 AVX-512 可用，AVX2 比较/移动掩码可能仍然是许多用例的最佳选择。 Intel CPU 仅在端口 5 上运行

vpcmpb

，与

kmov k, mem/reg

竞争。

kmov reg, k

在端口 0 上运行，与

kand

相同。

kmov mem, k

是一个纯存储，没有与负载不同的 ALU。 (https://uops.info/)

如果您很少需要将其作为不使用 512 位向量的较大程序的一部分，那么您可能也希望在这里避免使用它们。 AVX-512

kunpckdq

可以在一条指令中连接两个 32 位掩码，但它仅在 Intel 上的端口 5 上运行，并且您仍然需要将掩码数据输入或从

k

寄存器中取出。

---

没有 SVE 的 AArch64 有点不明显，最佳选择可能取决于您想对结果做什么。例如使用

shrn

进行比较并将字节缩小为半字节是有效的，以获得每个元素 4 位的掩码。 （这是在 16 字节向量中查找匹配位置的通常起点，例如作为

strlen

或

memchr

的一部分：将 /

shrn

与 4 和

fmov

与 64- 进行比较

clz

或

rbit/clz

的位整数寄存器。

因此，如果您的位图很小，将它们解压存储为半字节（每个元素 4 个掩码位）甚至字节（每个元素 8 位）将使这些操作更加高效。但如果它们（和数据）不适合 L1d 缓存，则值得花费更多周期来解包/重新打包为 1 位掩码元素。

我看到你的掩码只是一个

u64

，因此将掩码数据保存在 SIMD 向量中可能不会令人望而却步......但是对于诸如最低集之类的东西的位扫描对于标量掩码来说更有效。 如果您只需要测试任何非零值，您可以将一些掩码组合在一起并与零进行比较，然后 shrn/fmov 和

tst

该整数掩码以查看是否有任何掩码元素非零。