Rust 获取 SIMD 向量中真实字节的索引



我想比较两个 16 字节的向量并获取每个匹配的索引。一个小例子来说明我想要什么:

fn get_matching_idx(arr1: &[u8], arr2: &[u8]) {
let vec1 = u8x16::load_aligned(arr1);    
let vec2 = u8x16::load_aligned(arr2);
let matches = vec1.eq(vec2);
for i in 0..16 {
if matches.extract_unchecked(i) {
// Do something with the index
}
}
}

理想情况下,我只想为设置的索引"做点什么",而不是检查每一个索引(匹配数量会很少)。

有没有办法使用内部函数获取匹配索引,而不是遍历整个向量?例如,使用 gcc,我可以使用 _mm_movemask_epi8 对向量进行位打包,然后重复应用__builtin_clz来获取第一个设置位的索引(这对于我会拥有的稀疏数字来说性能更高)。或者,我可以有一个查找表,它为我的位打包整数中的每个半字节做了正确的事情(例如,这里的第一个答案)。

锈蚀中是否有这些指令的等效项?

我正在为英特尔 x86-64 处理器进行编译,不需要跨平台支持。

注意:我更喜欢原生(安全)锈蚀的解决方案,但这不是硬性要求。我可以写不安全的锈迹,甚至使用某种 FFI 链接到上述方法。

std::arch包含一组详尽的内部操作。这可以使用core::archstd::simd完成,如下所示:

use std::arch::x86_64::{self, __m128i};
use std::simd::{u8x16, FromBits};
unsafe fn get_matching_idx(arr1: &[u8], arr2: &[u8]) -> u32 {
let vec1 = __m128i::from_bits(u8x16::load_aligned_unchecked(arr1));
let vec2 = __m128i::from_bits(u8x16::load_aligned_unchecked(arr2));
return x86_64::_mm_movemask_epi8(x86_64::_mm_cmpeq_epi8(vec1, vec2)) as u32;
}
fn main() {
// let arr1 = ...
// let arr2 = ...
unsafe {
let mut mask = get_matching_idx(arr1, arr2);
}
let mut delta_i = 0;
// This assumes a little endian machine (note it counts trailing 0s)
while group_mask > 0 {
let tz = x86_64::_mm_tzcnt_32(mask);
let i = tz + delta_i;
// Do something...
group_mask >>= tz + 1;
delta_i += tz + 1;
}
}

最新更新