SIMD环境下带掩码模式匹配最优算法及Rust实现优化问询

需求描述 #

带掩码的模式匹配逻辑规则：对每个位置i，若mask[i]为1，则校验pat[i]与data[i]是否相等；若mask[i]为0，则直接判定为匹配。示例如下：

pat:    [1, 2, 3, 4]
data:   [1, 1, 3, 3]
mask:   [1, 0, 1, 1] # 1表示需校验相等，0表示无需校验
result: [1, 1, 1, 0] # 1表示匹配，0表示不匹配

原生算法逻辑 #

# 三元表达式版本
for i in range(len(pat)):
    result[i] = 1 if not mask[i] else (1 if pat[i] == data[i] else 0)

# 逻辑或版本
for i in range(len(pat)):
    result[i] = 1 if (not mask[i]) or (pat[i] == data[i]) else 0

现有Rust实现 #

1. 返回匹配掩码的Portable 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),
        )
}

2. 带运行时检测的多版本函数（判断是否全匹配） #

#[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()
}

架构专属最优实现建议 #

x86/x86_64系列 #

AVX512BW

利用AVX512专用掩码寄存器直接操作，无需将bit掩码扩展为字节：

1. 将mask加载到掩码寄存器k0；
2. 调用vpcmpb k1, zmm_data, zmm_pattern, 0x00得到逐字节相等的掩码k1；
3. 用knot k2, k0反转原掩码（标记无需校验的位置）；
4. 用kor k1, k1, k2合并两个掩码，得到最终匹配结果；
5. 调用ktest k1, k1判断是否所有位都为1（全匹配）。

对应Rust实现示例：

#[target_feature(enable = "avx512bw")]
unsafe fn simd_match_avx512bw<const N: usize>(pattern: [u8; N], mask: u64, data: [u8; N]) -> bool {
    let zmm_pat = _mm512_loadu_si512(pattern.as_ptr() as *const _);
    let zmm_data = _mm512_loadu_si512(data.as_ptr() as *const _);
    let k_mask = _mm512_loadu_epi64(&mask as *const _);
    let k_eq = _mm512_cmpeq_epi8_mask(zmm_data, zmm_pat);
    let k_not_mask = _mm512_knot(k_mask);
    let k_result = _mm512_kor(k_eq, k_not_mask);
    _mm512_ktest_mask(k_result, k_result) == 0
}

AVX2/SSE4.2/SSE2

这类指令集无专用掩码寄存器，需先将bit掩码扩展为逐字节掩码（mask第i位为1时，对应字节为0xFF，否则为0x00）：

1. 生成逐字节掩码数组；
2. 用_mm256_cmpeq_epi8（AVX2）或_mm_cmpeq_epi8（SSE）计算数据与模式的相等掩码；
3. 用_mm256_andnot_si256反转掩码字节，再与相等掩码做或运算得到最终结果；
4. 用_mm256_testz_si256（AVX2）或_mm_testz_si128（SSE）判断是否全匹配。

ARM/AArch64系列 #

AArch64 NEON

1. 将bit掩码扩展为逐字节掩码；
2. 调用vceqq_u8计算数据与模式的相等掩码；
3. 用vmvn_u8反转掩码字节，再与相等掩码做或运算；
4. 调用vminv_u8取结果中的最小字节，若为0xFF则表示全匹配。

AArch64 SVE/SVE2

利用SVE的谓词寄存器特性：

1. 将mask加载为谓词p2；
2. 调用cmpeq p0.b, z0.b, z1.b得到相等谓词；
3. 用not p1.b, p2.b反转原谓词；
4. 调用orr p0.b, p0.b, p1.b合并谓词；
5. 调用ptest p0.b, p0.b判断是否全匹配。

ARM NEON/VFP

逻辑与AArch64 NEON一致，针对32/64位向量使用对应指令（如vceq_u8、vorr_u8、vmvn_u8），通过vmin_u8判断结果是否全为0xFF。

多版本函数选型分析 #

你列出的targets覆盖了主流架构的核心SIMD指令集，可做以下优化调整：

• x86/x86_64：保留现有targets即可，sse2是基础指令集，avx/avx2/avx512bw为递进式高级指令集，无冗余。
• ARM：VFP2/VFP3对整数SIMD支持有限，若无需兼容极老设备，可移除arm+vfp2/arm+vfp3，优先保留arm+neon和arm+vfp4。
• AArch64：SVE/SVE2仅在ARMv8.2-A及以上版本支持，若目标环境包含这类设备则保留；若面向普通移动设备，aarch64+neon已足够覆盖大部分场景。

另外，建议为每个target编写架构专属实现，而非复用portable SIMD代码，这样能充分利用硬件特性，获得最优性能。

内容的提问来源于stack exchange，提问作者Steve Fan