Rust 如何优化小型固定大小数组中的搜索

发布于04月05日

我想找到一个16字节数组中字节的第一次出现.如果我编写了一个朴素的版本(使用迭代器或手动循环)，rustc看起来不会向量化(https://godbolt.org/z/fbKfvxTdv).

我对优化解决方案的模糊 idea 是

广播x到某个向量寄存器
把它和y的每一个元素进行比较
收集16个生成的布尔到u16
用count_leading_zeros做一些事情来找出哪个匹配.

相反，对于手动循环，它只是通过16次比较和跳转完全展开，而对于基于迭代器的版本，它似乎根本没有利用已知的固定长度(正好是16

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
}

推荐答案

每天晚上都有Portable SIMD module个，这是非常方便的.

请在下面找到实现你建议的算法的try ，以及编译器资源管理器中的反汇编版本，以及一个基准测试显示，至少在我的计算机上，你建议的SIMD版本是有益的.

#![feature(test)] // cargo bench
#![feature(portable_simd)]

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
}

pub fn find_first_simd(
    x: u8,
    y: &[u8; 16],
) -> Option<usize> {
    use std::simd::cmp::SimdPartialEq;
    use std::simd::u8x16;
    let x = u8x16::splat(x);
    let y = u8x16::from_array(*y);
    let e = x.simd_eq(y);
    e.first_set()
}
/*
in https://godbolt.org/
rustc nightly with options
  -C opt-level=3 -C target-feature=+avx -C target-feature=+avx2
gives

example::find_first_simd::hd89fd3b592dbfe03:
        vmovd   xmm0, edi
        vpbroadcastb    xmm0, xmm0
        vpcmpeqb        xmm0, xmm0, xmmword ptr [rsi]
        vpmovmskb       ecx, xmm0
        xor     eax, eax
        test    ecx, ecx
        setne   al
        rep       bsf edx, ecx
        ret
*/

fn main() {
    let y = [0, 1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1];
    for x in [0, 4, 9] {
        println!("first({}) {:?}", x, find_first(x, &y));
        println!("first_manual({}) {:?}", x, find_first_manual(x, &y));
        println!("first_simd({}) {:?}", x, find_first_simd(x, &y));
    }
}
/*
first(0) Some(0)
first_manual(0) Some(0)
first_simd(0) Some(0)
first(4) Some(4)
first_manual(4) Some(4)
first_simd(4) Some(4)
first(9) None
first_manual(9) None
first_simd(9) None
*/

#[cfg(test)]
mod tests {
    use super::*;
    extern crate test; // not in Cargo.toml

    const REPEAT: usize = 1000;

    #[bench]
    fn bench_find_first_simd(b: &mut test::Bencher) {
        b.iter(|| {
            for _ in 0..REPEAT {
                let y = [0, 1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1];
                for x in [0, 4, 9] {
                    test::black_box(find_first_simd(
                        test::black_box(x),
                        test::black_box(&y),
                    ));
                }
            }
        });
    }

    #[bench]
    fn bench_find_first(b: &mut test::Bencher) {
        b.iter(|| {
            for _ in 0..REPEAT {
                let y = [0, 1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1];
                for x in [0, 4, 9] {
                    test::black_box(find_first(
                        test::black_box(x),
                        test::black_box(&y),
                    ));
                }
            }
        });
    }

    #[bench]
    fn bench_find_first_manual(b: &mut test::Bencher) {
        b.iter(|| {
            for _ in 0..REPEAT {
                let y = [0, 1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1];
                for x in [0, 4, 9] {
                    test::black_box(find_first_manual(
                        test::black_box(x),
                        test::black_box(&y),
                    ));
                }
            }
        });
    }
}
/*
test tests::bench_find_first        ... bench:       4,282 ns/iter (+/- 74)
test tests::bench_find_first_manual ... bench:       4,246 ns/iter (+/- 45)
test tests::bench_find_first_simd   ... bench:       2,716 ns/iter (+/- 45)
*/