シクシク素数列(Rust) を高速化する

シクシク素数列 Advent Calendar 2018 (Rust) をあげたところ優しい方たちからアドバイスをもらったので、実行速度に注目して色々やっていく。

解こうとしている問題そのものはシクシク素数列 Advent Calendar 2018をご覧ください。

途中ごちゃごちゃやっているので、最終結果だけみたい場合は最後までスクロールしてください。

準備

測定環境

• OS: macOS High Sierra

$ system_profiler SPHardwareDataType      
Hardware:

    Hardware Overview:

      Model Name: iMac
      Model Identifier: iMac16,2
      Processor Name: Intel Core i7
      Processor Speed: 3.3 GHz
      Number of Processors: 1
      Total Number of Cores: 4
      L2 Cache (per Core): 256 KB
      L3 Cache: 6 MB
      Memory: 16 GB

Rust

$ cargo version
cargo 1.32.0-nightly (5e85ba14a 2018-12-02)

（nightlyなのに大きな意味はないです。今回のコードは stable でも動くはず）

測定方法

• 適切に大きなシクシク素数を見つけるのにかかる時間を計測する
• cargo build --release でリリースビルドを用いる
• 出力は /dev/null に捨てる

実際のコマンド例

$ cargo build --release
$ echo 10000 | /usr/bin/time ./target/release/qiita4949 > /dev/null

チューニング開始

計算に用いるプリミティブ型を変更

最初の記事では Rust のプリミティブ型としてもっとも大きな整数を表すことができるからという意味で u128 (128bit符号無し整数)を用いた。
しかし、今回の問題設定ではそんなに大きな数は扱わないので実際問題として i32 (32bit符号付整数)で十分。
それぞれの型を用いたときの実行速度を測定しておく。

それぞれ1回ずつのみの計測。単位は秒。

計算に用いるプリミティブ型	1万	2万	4万
i32	2.37	10.14	38.01
u32	2.99	12.86	47.57
i64	8.05	35.04	127.08
u64	7.58	32.11	119.72
i128	22.40	96.38
u128	20.51	87.74

思っていた以上に計算速度が違った。
bit数が多いと計算に時間がかかるのはなんかそんな気もする。
32bitだと符号付の方が早くて、64bit以上だと符号無しの方が早いのはよくわからない。アセンブラとかCPUの演算器とかに詳しい人に聞いてみたい。

とりあえず、以後は i32 を用いる。

素数判定を sqrt で打ち切る

ループの上限を sqrt であらかじめ計算してもいいんだが、nが平方数のときに計算誤差をうまく処理するのがめんどくさいので（アドホックに+2とかすれば防げる問題ではある）、 以下のように処理する。

type Integer = i32;

fn is_prime(n: Integer) -> bool {
    if n <= 1 {
        return false;
    }
    for i in 2..n {
        if n % i == 0 {
            return false;
        }
        if i * i > n {
            return true;
        }
    }
    true
}

#[test]
fn test_prime() {
    assert_eq!(is_prime(0), false);
    assert_eq!(is_prime(1), false);
    assert_eq!(is_prime(2), true);
    assert_eq!(is_prime(9), false);
    assert_eq!(is_prime(57), false);
    assert_eq!(is_prime(97), true);
}

ユニットテストも通ったので、早速1万個のシクシク素数を求めてもらう。

$ echo 10000 | /usr/bin/time ./target/release/qiita4949 > /dev/null 
        0.01 real         0.01 user         0.00 sys

あ、ハイって感じですね。
これだとわからないので100万番目のシクシク素数までやってみます。

$ echo 1000000 | /usr/bin/time ./target/release/qiita4949 > /dev/null
       11.31 real        11.28 user         0.03 sys

いい線いってる。

入手した素数を貯めておく

個人的によく使う素数判定アルゴリズム。自作。
発想はエラトステネスの篩と大きくは変わらないと思うが、　コードが簡単なのでささっと書くときに使いがち。

type Integer = i32;

struct PrimeSeries {
    primes: Vec<Integer>,
}

impl PrimeSeries {
    fn new() -> PrimeSeries {
        PrimeSeries {
            primes: vec![2, 3, 5],
        }
    }

    fn expand(&mut self, max2: Integer) {
        for i in (self.primes.last().unwrap() + 2..max2).step_by(2) {
            if i * i > max2 {
                break;
            }
            if self.is_prime(i) {
                self.primes.push(i);
            }
        }
    }

    /// 現在の自前の素数列だけから素数判定を行う
    /// 適切に ``expand`` しておかないとバグる
    /// 特にこだわりがないなら ``determine_prime`` を使うべき
    fn is_prime(&self, n: Integer) -> bool {
        if n <= 1 {
            return false;
        }
        for p in &self.primes {
            if p * p > n {
                return true;
            }
            if n % p == 0 {
                return false;
            }
        }
        true
    }

    fn determine_prime(&mut self, n: Integer) -> bool {
        self.expand(n);
        self.is_prime(n)
    }
}

#[test]
fn test_prime() {
    let mut ps = PrimeSeries::new();
    assert_eq!(ps.determine_prime(0), false);
    assert_eq!(ps.determine_prime(1), false);
    assert_eq!(ps.determine_prime(2), true);
    assert_eq!(ps.determine_prime(9), false);
    assert_eq!(ps.determine_prime(57), false);
    assert_eq!(ps.determine_prime(97), true);
    assert_eq!(ps.determine_prime(9999991), true);
}

計測

$ echo 1000000 | /usr/bin/time ./target/release/qiita4949 > /dev/null                                                                                         
        2.73 real         2.69 user         0.02 sys

まあそんなもんですよね。

エラトステネスの篩を借りてくる

これを使う。
コード中のコメントもコピーしてきたものなので僕のものではない。

ただし以下の変更を加えた

• Rust　の今のバージョンでは step_byがあるので利用する
• 条件を揃えるため、サイズが見積もりにくい、などの理由で初期キャパシティはデフォルトを利用

type Integer = i32;

// FNVハッシュはキーのビット数が小さい時はSipHashよりも高速。
use fnv::FnvHashMap;

fn primes() -> Box<Iterator<Item = Integer>> {
    let mut multiples = FnvHashMap::with_hasher(Default::default());;
    let iter = (3..).step_by(2).filter_map(move |i| {
        let (prime_or_none, factor) = match multiples.remove(&i) {
            Some(f) => (None, f),
            None => (Some(i), i * 2),
        };

        (1..)
            .map(|j| i + j * factor)
            .skip_while(|m| multiples.contains_key(m))
            .next()
            .map(|m| multiples.insert(m, factor));

        prime_or_none // filter_map()は次の素数が見つかるまでリトライする。
    });
    // 最初に2を返してから、他の素数を返す。
    Box::new((2..).take(1).chain(iter))
}

fn solve_4949(n: usize) -> Vec<Integer> {
    let mut ans = Vec::new();
    for p in primes() {
        if is_4949(p) {
            ans.push(p);
            if ans.len() >= n {
                break;
            }
        }
    }
    ans
}

#[test]
fn test_primes() {
    let ans = vec![2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31];
    for (p1, p2) in primes().zip(ans) {
        assert_eq!(p1, p2);
    }
}

計測

$ echo 1000000 | /usr/bin/time ./target/release/qiita4949 > /dev/null
        1.97 real         1.92 user         0.03 sys

だいぶ早くなった。

最終まとめ

初期版とはちょっと比較できないくらい早くなった。
それこそ桁がいくつも違う。

use std::io::stdin;

// FNVハッシュはキーのビット数が小さい時はSipHashよりも高速。
use fnv::FnvHashMap;

type Integer = i32;

/// 参考: https://qiita.com/termoshtt/items/9ddcc9d6c56c3538558d#comment-bc87fc360e280ab70bff
fn primes() -> Box<Iterator<Item = Integer>> {
    let mut multiples = FnvHashMap::with_hasher(Default::default());;
    let iter = (3..).step_by(2).filter_map(move |i| {
        let (prime_or_none, factor) = match multiples.remove(&i) {
            Some(f) => (None, f),
            None => (Some(i), i * 2),
        };

        (1..)
            .map(|j| i + j * factor)
            .skip_while(|m| multiples.contains_key(m))
            .next()
            .map(|m| multiples.insert(m, factor));

        prime_or_none // filter_map()は次の素数が見つかるまでリトライする。
    });
    // 最初に2を返してから、他の素数を返す。
    Box::new((2..).take(1).chain(iter))
}

fn is_4949(mut n: Integer) -> bool {
    loop {
        if n == 0 {
            return false;
        }
        match n % 10 {
            4 | 9 => return true,
            _ => n /= 10,
        }
    }
}

fn solve_4949(n: usize) -> Vec<Integer> {
    let mut ans = Vec::new();
    for p in primes() {
        if is_4949(p) {
            ans.push(p);
            if ans.len() >= n {
                break;
            }
        }
    }
    ans
}

fn main() {
    let mut buf = String::new();
    stdin().read_line(&mut buf).unwrap();
    let n = buf.trim().parse().expect("fail parse as integer");
    let ans = solve_4949(n);
    println!(
        "{}",
        ans.iter()
            .map(|n| n.to_string())
            .collect::<Vec<String>>()
            .join(","),
    );
}

#[test]
fn test_primes() {
    let ans = vec![2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31];
    for (p1, p2) in primes().zip(ans) {
        assert_eq!(p1, p2);
    }
}

#[test]
fn test_4949() {
    assert_eq!(is_4949(0), false);
    assert_eq!(is_4949(1), false);
    assert_eq!(is_4949(4), true);
    assert_eq!(is_4949(9), true);
    assert_eq!(is_4949(87318112), false);
    assert_eq!(is_4949(987318112), true);
}