performance - 优化 rust N 素数生成器

Question

玩rust我实现的迭代器和 lambda N素数生成器 ( original playground/improved code. )。
与我在同一个盒子上更精通的其他语言的等效功能代码相比，它的性能非常差(慢 10 倍)。
一些数字:

test calc_primes ... bench: 7,754,795 ns/iter (+/- 1,887,591)

#[macro_use]
extern crate bencher;
use bencher::Bencher; 
use nth_prime::*;                                                                                                   
fn calc_primes(b: &mut Bencher) {
   b.iter(|| primes(10000)); 
} 

benchmark_group!(benches, calc_primes); 
benchmark_main!(benches);                           

我希望 Rust 能胜过它，但我很难消除以下代码中的缺陷。因此，这里呼吁帮助优化此特定实现。

• 我不喜欢那些Vec分配，有些肯定不需要...
• 可能到处传递更多的引用
• 其他实现使用逻辑 OR关于 fvec bitarray这是 super 快。下面 rust 中的逻辑在 fold 中表达了相同的方法。调用，但这是使用 Vec<bool>与等效的位数组相比，这一定很慢...
• 使用不同的数据类型代替 Vec<bool> , 不介意分配大量字节数组并使用位移等操作，但无法想象将它们放在一起......

pub fn primes(n: u64) -> Vec<usize> {                                                                                                                                             
    if n < 4 {                                                                                                                                                                    
        vec![2]                                                                                                                                                                   
    } else {                                                                                                                                                                      
        base(primes((n as f64).sqrt().ceil() as u64), n)                                                                                                                          
    }                                                                                                                                                                             
}                                                                                                                                                                                 

fn base(r: Vec<usize>, p: u64) -> Vec<usize> {                                                                                                                                    
    let fvec = (0..p).map(|_| false).collect::<Vec<bool>>();                                                                                                                      
    let z = r                                                                                                                                                                     
        .iter()                                                                                                                                                                   
        .map(|&x| (0..x).map(|y| y > 0).collect::<Vec<bool>>())                                                                                                                   
        .map(|x| {                                                                                                                                                                
            (0..p)                                                                                                                                                                
                .map(|n| !x[(n % (x.len() as u64)) as usize])                                                                                                                     
                .collect::<Vec<bool>>()                                                                                                                                           
        })                                                                                                                                                                        
        .fold(fvec, |acc, x| {                                                                                                                                                    
            acc.iter().zip(x).map(|(a, b)| *a || b).collect()                                                                                                                     
        });                                                                                                                                                                       
    let yy: Vec<usize> = z                                                                                                                                                        
        .iter()                                                                                                                                                                   
        .enumerate()                                                                                                                                                              
        .filter_map(|(i, x)| if !*x { Some(i) } else { None })                                                                                                                    
        .skip(1)                                                                                                                                                                  
        .collect();
    r.iter().chain(yy.iter()).map(|x| *x).collect()                                                                                                                               
}                                                                                          

关于素数生成算法，在this thread中有出色的答案，所以这不是关于优化算法的问题，而是这个没有外部 crate 的特定实现(这是一个学习练习)。

编辑:

经过评论部分的一些提示后，对其进行了显着优化:

test calc_primes       ... bench:   4,776,400 ns/iter (+/- 1,427,534)
test calc_primes_sieve ... bench:     644,220 ns/iter (+/- 1,655,581)
test calc_primes_2     ... bench:     268,598 ns/iter (+/- 46,440)

1. 替换为“低效”Vec<bool> @ fold具有惯用迭代器方式的位图:step_by
2. 放弃不可变类型以支持状态(有没有办法保持不可变方法？)

fn base2(head: Vec<u64>, p: u64) -> Vec<u64> {                                                                                                                                    
    let mut fvec = (0..p).map(|_| false).collect::<Vec<bool>>();                                                                                                                  
    head.iter()                                                                                                                                                                   
        .map(|&x| {                                                                                                                                                               
            (0..p)                                                                                                                                                                
                .step_by(x as usize)                                                                                                                                              
                .for_each(|y| fvec[y as usize] = true)                                                                                                                            
        })                                                                                                                                                                        
        .for_each(drop);                                                                                                                                                          
    let tail: Vec<u64> = fvec.iter().enumerate()
        .filter_map(|(i, x)| if !*x { Some(i as u64) } else { None })                                                                     
        .skip(1)                                              
        .collect();                                     
    head.iter().chain(tail.iter()).cloned().collect()                                                                                                                             
}                                                         

编辑 2:

作为引用，这里是“其他”语言的实现:

q)p:{$[x<4; enlist 2; r, 1_ where not any x#'not til each r:p ceiling sqrt x]}



q)\ts do[1000;p 10000]
474 443088  
=
474,000 ns/iter 

对于 64 位 q v4.0:

q)\t do[1000;p 10000]
273  
=
273,000 ns/iter 

对于解释性语言来说还不错。

编辑 3

为来自答案的新实现添加了更多基准。
(按时间顺序出现)
能够在没有skip(1)的情况下剃掉一点点和差异 fvec分配 New playground code.

test ssqrt_p_1__vec_bool_n_mod   ... bench:   6,838,150 ns/iter (+/- 627,389)
test sieve_p_2__mut_step_by      ... bench:     367,229 ns/iter (+/- 38,279)
test ssqrt_p_3__mut_step_by      ... bench:     266,189 ns/iter (+/- 56,215)
test ssqrt_p_4__push_step_by_mod ... bench:   1,255,206 ns/iter (+/- 262,107)
test sieve_p_5__for_loop_mut     ... bench:     441,397 ns/iter (+/- 47,077)
test ssqrt_p_6__no_skip          ... bench:     176,186 ns/iter (+/- 24,765)

StdDev 现在起着重要作用......运行:

taskset -c 0 cargo bench --jobs 1

fn base3(head: Vec<u64>, p: u64) -> Vec<u64> {                                                                                                                                    
    let mut fvec = vec![false; p as usize]; fvec[1] = true;                                                                                                                       
    head.iter()                                                                                                                                                                   
        .map(|&x| {                                                                                                                                                               
            (0..p)                                                                                                                                                                
                .step_by(x as usize)                                                                                                                                              
                .for_each(|y| fvec[y as usize] = true)                                                                                                                            
        })                                                                                                                                                                        
        .for_each(drop);                                                                                                                                                          
    let tail: Vec<u64> = fvec                                                                                                                                                     
        .iter()                                                                                                                                                                   
        .enumerate()                                                                                                                                                              
        .filter_map(|(i, x)| if !*x { Some(i as u64) } else { None })                                                                                                             
        .collect();                                                                                                                                                               
    head.iter().chain(tail.iter()).cloned().collect()                                                                                                                             
} 

编辑 5

令人惊讶的是下面no_chain比 base3:no_skip 慢版本。我猜这就是编译器的魔法 (递归函数中的尾调用优化 ???)。

fn base4(head: Vec<u64>, p: u64) -> Vec<u64> {                                                                                                                                    
    let mut fvec = vec![false; p as usize]; fvec[1] = true;                                                                                                                       
    head.iter()                                                                                                                                                                   
        .map(|&x| {                                                                                                                                                                           
            (0..p)                                                                                                                                                                
                .step_by(x as usize)                                                                                                                                              
                .for_each(|y| fvec[y as usize] = true);                                                                                                                           
            fvec[x as usize] = false;                                                                                                                                             
        })                                                                                                                                                                        
        .for_each(drop);                                                                                                                                                          
    fvec.iter().enumerate()                                                                                                                                                       
        .filter_map(|(i, x)| if !*x { Some(i as u64) } else { None })                                                                                                             
        .collect()                                                                                                                                                                
}                                              

Answer 1

你的代码有一些非常奇怪的开销，比如

(0..x).map(|y| y > 0).collect::<Vec<bool>>()

你只能像 x[(n % (x.len() as u64)) as usize]... 又名。作为编写 n % x == 0 的一种非常缓慢的方式。

这是您的代码的一个几乎直接的变体，我只是删除了您添加的这些开销，并在有意义的地方内联代码。

pub fn primes(n: usize) -> Vec<usize> {
    if n < 4 {
        vec![2]
    } else {
        let mut r = primes((n as f64).sqrt().ceil() as usize);
        let mut fvec: Vec<_> = (0..n).map(|i| r.iter().any(|x| i % x == 0)).collect();
        fvec[1] = true;

        for (i, x) in fvec.into_iter().enumerate() {
            if !x { r.push(i) }
        }
        r
    }
}

最大的实际变化是将 skip(1) 替换为 fvec[1] = true，因为这样会简单得多。

当然，实际的埃拉托色尼筛法更简单，开销更低，时间复杂度更好。

pub fn primes(n: usize) -> Vec<usize> {
    let mut primes = Vec::new();
    let mut sieve = vec![true; n];

    for p in 2..n {
        if sieve[p] {
            primes.push(p);
            let mut i = p * p;
            while i < n {
                sieve[i] = false;
                i += p;
            }
        }
    }

    primes
}

这是一个略微优化的筛子，避免了不必要的工作和内存使用。

pub fn primes6(mut n: usize) -> Vec<usize> {
    let mut primes = vec![2];
    n >>= 1;

    // 0  1  2  3  4
    // 2, 3, 5, 7, 9, ...
    let mut sieve = vec![true; n];

    let end = (n as f64).sqrt().ceil() as usize;
    for p in 1..end {
        if sieve[p] {
            let prime = p * 2 + 1;
            primes.push(prime);

            let mut i = ((prime * prime) - 1) / 2;
            while i < n {
                sieve[i] = false;
                i += prime; // skip by 2·prime
            }
        }
    }

    for p in end..n {
        if sieve[p] {
            let prime = p * 2 + 1;
            primes.push(prime);
        }
    }

    primes
}