【Rust】マルチスレッドで素数を数えて落ち着く

こんにちは！
ここ半年くらいAIにコードを書いてもらっているのですが、便利な反面、徐々にコードを書く筋力が落ちてるな、と感じています

そこで、自主トレがてら、Rustでアルゴリズムを書いて遊んでいます

ルール

1. 自力でコードを書く
2. Codex先生にレビューしてもらう
3. 直す
4. Codex先生が良いと言うまで続ける

言語選択

言語はRustを選択しました

• 書くために考えなければいけないことが多い
• 変数ひとつ定義するにしても、スタックに置くのか、ヒープアロケーションするのか、考えないといけない
• 関数を定義するなら、引数は所有権ごともらうのか、参照でもらうのか、ミュータブルなのかイミュータブルなのか、考えないといけない
• マルチスレッドならCPU側の仕様（メモリオーダリングなど）も考慮しないといけない

コードを書く筋肉を取り戻すには良い言語なのではという判断
Rustで正しく実装するのは本当に難しい

テーマ

与えられた自然数の配列に含まれる素数の数をカウントする
マルチスレッド対応で、実行するマシンのコア数に応じて並行処理する

実装

use std::thread;

/// 与えられた `numbers` の中に含まれる素数の個数を並列に数える
/// スレッド数は利用可能な並列度と入力長に基づき自動決定する
pub fn prime_explorer(numbers: &[u64]) -> u64 {
    let len = numbers.len();
    if len == 0 {
        return 0;
    }

    // 利用可能なスレッド数を取得し、入力長を上限にして空チャンクを作らないようにスレッド数を決定
    let threads = thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(1)
        .min(len);

    if threads == 1 {
        return prime_counter(numbers);
    }

    // 切り上げ割り算でチャンクサイズを決め、すべての要素をどこかに含める
    let chunk_size = (len + threads - 1) / threads;

    thread::scope(|s| {
        // 先に全チャンクを spawn してハンドルを集め、その後 join して並列に実行する
        let handles: Vec<_> = numbers
            .chunks(chunk_size)
            .map(|chunk| s.spawn(move || prime_counter(chunk)))
            .collect();

        handles.into_iter().map(|h| h.join().unwrap()).sum()
    })
}

/// スライス内の素数の個数を数える
fn prime_counter(numbers: &[u64]) -> u64 {
    numbers
        .iter()
        .copied()
        .fold(0u64, |acc, n| acc + is_prime(n) as u64)
}

/// 与えられた数が素数かを判定する
fn is_prime(number: u64) -> bool {
    if number < 2 {
        return false;
    }
    if number == 2 {
        return true;
    }
    if number % 2 == 0 {
        return false;
    }

    // 浮動小数の sqrt を避け、i*i <= number を i <= number/i で判定して試し割り
    let mut i = 3;
    while i <= number / i {
        if number % i == 0 {
            return false;
        }
        i += 2;
    }
    true
}

解説

今回の実装では、

• CPU バウンドな処理（素数判定）
• 共有状態を持たない独立な計算
• 入力サイズが事前に分かっている

という性質を活かし、「データを分割して並列に処理し、最後に合計する」という実装にしています

全体構成

処理の流れは次の通りです

1. 入力スライスの長さを確認し、空なら即終了
2. 利用可能な CPU コア数からスレッド数を決定
3. 入力をスレッド数に応じてチャンクに分割
4. 各チャンクを別スレッドで素数カウント
5. 各スレッドの結果を join して合計

---

スレッド数の決定

let threads = thread::available_parallelism()
    .map(|n| n.get())
    .unwrap_or(1)
    .min(len);

• available_parallelism() は、実行環境で同時に実行可能なスレッド数（論理コア数）を返します
• 入力長より多いスレッドを作っても余ってしまうため、min(len) で上限をかけています

---

チャンク分割

let chunk_size = (len + threads - 1) / threads;

ここでは 切り上げ除算を使ってチャンクサイズを計算しています

• すべての要素が必ずどこかのチャンクに含まれる
• 最後のチャンクだけ小さくなる可能性はあるが、許容

---

thread::scope を使った並行処理

thread::scope(|s| {
    let handles: Vec<_> = numbers
        .chunks(chunk_size)
        .map(|chunk| s.spawn(move || prime_counter(chunk)))
        .collect();

    handles.into_iter().map(|h| h.join().unwrap()).sum()
})

マルチスレッド部分です

• thread::scope を使うことで、numbers の 参照（borrow）をそのままスレッドに渡せる
• Arc や Mutex を使わずに並行処理ができる
• 各スレッドは独立したスライスを処理するため、共有状態・ロックが一切ない

---

素数カウント処理

fn prime_counter(numbers: &[u64]) -> u64 {
    numbers
        .iter()
        .copied()
        .fold(0u64, |acc, n| acc + is_prime(n) as u64)
}

• fold を使い、素数なら 1、そうでなければ 0 を加算
• bool as u64 による変換で、分岐を最小限に

単一スレッドでもそのまま使える、純粋関数として切り出しています
boolをu64に変換すると0と1になるやつ、低レイヤーを感じる...！
※ 面白そうなのでfoldにしていますが、パフォーマンス出すならmapとかfilterとか使った方がいいかも？？ このへん検証すると楽しいかもしれんね

---

素数判定アルゴリズム

fn is_prime(number: u64) -> bool {

素数判定は、みんな大好き試し割り法をベースにしています

• 2 未満は素数ではない
• 2 は素数
• 偶数は除外
• 奇数のみを √n まで試し割り

whileのところの条件式がポイント

while i <= number / i {

• i * i <= number の代わりに i <= number / i を使うことで整数オーバーフローの可能性を排除

このへんも低レイヤーを感じるぜ...！

まとめ

Rustめっちゃ難しいけど楽しい
趣味で遊ぶにはいい言語だ...