如何优化16字节固定数组的首匹配字节搜索？Rust编译器未向量化问题

问题背景 #

需要在长度固定为16的字节数组中查找某个字节的首次出现位置，但使用迭代器或手动循环的朴素实现时，rustc并未进行向量化优化：手动循环版本被完全展开为16次比较与跳转操作，迭代器版本也未利用数组长度固定的特性。

预期优化思路 #

• 将目标字节x广播至向量寄存器
• 与数组y的所有元素同时进行比较
• 将16个布尔比较结果打包为u16类型的掩码
• 利用位运算（如trailing_zeros）找到首个匹配的位置

现有朴素实现代码 #

pub fn find_first(x: u8, y: &[u8;16]) -> Option<usize> {
    y.iter().position(|w| *w == x)
}

pub fn find_first_manual(x: u8, y: &[u8;16]) -> Option<usize> {
    for i in 0..16 {
        if y[i] == x {
            return Some(i);
        }
    }
    None
}

手动实现向量化优化版本 #

针对x86_64架构，可以利用SIMD指令集（如SSE2）手动实现向量化查找，示例代码如下：

#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;

#[cfg(target_arch = "x86_64")]
pub fn find_first_simd(x: u8, y: &[u8; 16]) -> Option<usize> {
    unsafe {
        // 广播目标字节到128位向量寄存器
        let xmm_target = _mm_set1_epi8(x as i8);
        // 加载数组到向量寄存器
        let xmm_array = _mm_loadu_si128(y.as_ptr() as *const __m128i);
        // 逐字节比较，相等位置设为0xFF，否则0x00
        let cmp_result = _mm_cmpeq_epi8(xmm_target, xmm_array);
        // 将比较结果转换为u16掩码（每一位对应一个字节的比较结果）
        let mask = _mm_movemask_epi8(cmp_result) as u16;

        if mask == 0 {
            None
        } else {
            // 计算掩码中首个1的位置，即数组中首个匹配元素的索引
            Some(mask.trailing_zeros() as usize)
        }
    }
}

// 非x86_64架构下回退到手动循环实现
#[cfg(not(target_arch = "x86_64"))]
pub fn find_first_simd(x: u8, y: &[u8; 16]) -> Option<usize> {
    find_first_manual(x, y)
}

优化说明 #

这个版本通过SIMD指令一次性完成16个字节的并行比较，避免了循环展开带来的多次条件跳转。_mm_movemask_epi8将向量中的每个字节的符号位（比较相等时字节为0xFF，符号位为1）打包成整数掩码，再通过trailing_zeros快速定位到首个匹配的索引位置，效率远高于朴素实现。

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