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

Posted

技术标签:

【中文标题】Rust 获取 SIMD 向量中真实字节的索引【英文标题】:Rust get index of true bytes in SIMD vector 【发布时间】:2018-04-10 17:53:56 【问题描述】:

我想比较两个 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 以获取第一个设置位的索引(这对于我将拥有的稀疏数字更有效)。或者,我可以有一个查找表,它为我的位压缩整数中的每个半字节做正确的事情(例如,第一个答案 here)。

rust 中是否有与这些指令等效的指令?

我正在为 Intel x86-64 处理器编译,不需要跨平台支持。

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

【问题讨论】:

你可以在 Rust 中通过std::arch 使用相同的内在函数:doc.rust-lang.org/nightly/core/arch/x86_64/… --- 请注意,这是一个仅在夜间使用的 API,但计划很快就会稳定下来。如果您需要在稳定的 Rust 上执行此操作,那么最简单的方法可能是在 C 中纠正您的 SIMD 例程。 【参考方案1】:

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;
    

【讨论】:

为什么是core::arch 而不是std::arch?另外,为什么是_mm_tzcount_32 而不是u32::trailing_zeroes 看起来 _mmtzcount_32 只是调用汇编指令 tzcntl 而 u32::trailing_zeros 做得更多(但如果我错了,请随时纠正我):play.rust-lang.org/… 更新为使用 std::arch trailing_zeros 有一个 0 参数的检查;我不确定 0 参数的内在行为,我的经验是 gcc 内在函数在这方面指定 0 参数会导致未定义的行为。 @MatthieuM。 _mm_tzcnt_32(通过与tzcnt 指令等效)为零输入定义。 GCC 内置可能未指定为 0,因此它可以使用中断为 0 的指令(序列),例如最简单的 bsf 用法(当输入为零时不会写入其目的地)

以上是关于Rust 获取 SIMD 向量中真实字节的索引的主要内容,如果未能解决你的问题,请参考以下文章

使用 SIMD 对半字节的去交错向量

使用 F# 和 SIMD 搜索值索引

带字节分隔符的可变大小字段的快速 SIMD 提取

是否可以将字节解码为 UTF-8,将错误转换为 Rust 中的转义序列?

RGB 到 YCbCr 使用 SIMD 向量丢失一些数据

获取向量的字节大小[重复]