🏎️

Rust ndarray並列化 躓きながら3倍高速化

2024/06/25に公開

ndarrayを使用して計算を数値計算など行っているなら、イテレータをrayonによる並列化関数に置き換えるだけで並列処理を行うことができる!
リソースのロックとか(よくわからないけど)全く考えることなく並列化による高速化の恩恵を受けることができるのはすごい。
今回はじめて並列化に取り組んだ。
rayon(レーヨン)を使用すると「簡単に」並列化ができると書いてあったが、うまくいかなかったので記録。

結論

ndarray::parallel を使用する。
cargo.tomlに記載する依存関係で下記のようにfeaturesを記述することで使用できるようになる。

ndarray = {version = "0.15",features = ["rayon"]}

情報源

下記に書いてあった。英語弱者向けに、tomlのサンプルコード書いておいてほしかった。
featuresの指定についてよくわかっていなかった。

Module ndarray::parallel

迷った経緯

ndarrayの並列化方法を調べると、ndarray-parallelというクレートを使用するのだという情報が見つかり、しばらく試したが、できなかった。
現在は非推奨となっていた。一次情報を先に確認するの大事。
今は非推奨 ndarray-parallel

困っていたこと

コンパイルすると、並列化メソッドが存在しないといわれてしまう。

error[E0599]: the method `into_par_iter` exists for struct `AxisIterMut<'_, f64, Dim<[usize; 1]>>`, but its trait bounds were not satisfied

効果確認

並列化したら計算時間が1/3以下になった

# serial calc: 3.6764023s
# parallel calc : 1.1340838s
The arrays are equal.

並列化の効果確認コード

追記:下記コードでは、厳密に並列化の効果だけの検証にはなっていなかった。
forループでインデックスを回すのか、イテレータで処理をするのかでも速さが変わってくる。
厳密に並列化有無のみ比較するならpar_calc()におけるinto_par_iter()をinto_iter()に変えるとよい。

use ndarray::parallel::prelude::*;
///raion の技術習得
use ndarray::prelude::*;
use ndarray::Array2;
use ndarray::OwnedRepr;
use std::time::Instant; //パフォーマンス計測用

fn main() {
    let time_start = Instant::now(); // 各ループの開始時間を記録
    let res_serial = serial_calc();
    let time_duration = time_start.elapsed();
    println!("# serial calc: {:?}", time_duration);

    let time_start = Instant::now(); // 各ループの開始時間を記録
    let res_par = par_calc();
    let time_duration = time_start.elapsed();
    println!("# parallel calc : {:?}", time_duration);

    //結果一致確認
    let are_equal = res_serial == res_par;

    if are_equal {
        println!("The arrays are equal.");
    } else {
        println!("The arrays are not equal.");
    }
}

fn serial_calc()-> ndarray::ArrayBase<OwnedRepr<f64>, ndarray::Dim<[usize; 2]>> {
    let nx = 200;
    let ny = 200;
    let preheat_k = 100.0;
    let mut u: ArrayBase<ndarray::OwnedRepr<f64>, Dim<[usize; 2]>> =
        Array2::<f64>::from_elem((nx, ny), preheat_k);
    let mut un: ArrayBase<ndarray::OwnedRepr<f64>, Dim<[usize; 2]>> =
        Array2::<f64>::from_elem((nx, ny), preheat_k);

    for __ in 1..1000 {
        un.assign(&u);

        for i in 1..(nx - 1) {
            for j in 1..(ny - 1) {
                u[[i, j]] = un[[i, j]] + 1.0;
            }
        }
    }
    return u;
}

fn par_calc() -> ndarray::ArrayBase<OwnedRepr<f64>, ndarray::Dim<[usize; 2]>>{
    let nx = 200;
    let ny = 200;
    let preheat_k = 100.0;
    let mut u: ArrayBase<ndarray::OwnedRepr<f64>, Dim<[usize; 2]>> =
        Array2::<f64>::from_elem((nx, ny), preheat_k);
    let mut un: ArrayBase<ndarray::OwnedRepr<f64>, Dim<[usize; 2]>> =
        Array2::<f64>::from_elem((nx, ny), preheat_k);

    for __ in 1..1000 {
        un.assign(&u);

        let mut u_slice = u.slice_mut(s![1..(nx - 1), 1..(ny - 1)]);
        let un_slice = un.slice(s![1..(nx - 1), 1..(ny - 1)]);

        // 並列化された計算
        u_slice
            .axis_iter_mut(ndarray::Axis(0))
            .into_par_iter()
            .enumerate()
            .for_each(|(i, mut row)| {
                let un_row = un_slice.row(i);
                row.iter_mut().enumerate().for_each(|(j, elem)| {
                    *elem = un_row[j] + 1.0;
                });
            });
    }
    return u;
}

Discussion