rustのfor文をマルチスレッドにしたい

Question

私は下記の関数を並列処理に変更したいと考えています。

fn cal_v(c: &Array3<f32>, dt: f32) -> Array3<f32> {
  let mut v: Array3<f32> = Array::zeros(c.raw_dim());
  for time in 1..c.shape()[0] {
    let before = c.slice(s![time-1, .., ..]).to_owned();
    let now = c.slice(s![time, .., ..]).to_owned();
    let now_v = (before - now) / dt;
    v.slice_mut(s![time, .., ..]).assign(&now_v);
  }
  v
}

この関数のcは3次元の配列で1次元目が時間を表すindexで2次元目は粒子の番号を表すindexで3次元目はx,y,zを表す次元です。
個人的にマルチスレッドにする書き方を学びたいので、
for time in 1..c.shape()[0]の部分をマルチスレッドで行いたいと考えております。

どのような方法で書くのがよろしいでしょうか？

使用しているコンパイラのバージョンは、1.51.0
使用しているライブラリのバージョンは、ndarray(0.14.0)です。

Comments

[blackenedgold]

そのプログラムですとfor式内でbeforeを参照している（1つ前のイテレーションの出力に依存している）ためloop carried dependencyがあるように見えます。つまり、並列化できないように見えます。
また、 `assgin` 等綴り間違いもあります。コンパイルが通って正しい結果を返すプログラムを掲載頂けないでしょうか。

また、コンパイラのバージョンと使っているライブラリ（恐らくndarrayでしょうか）のバージョンもお願いします。

[tatsuya6502]

beforeとnowがゼロで初期化したように見えるvから値を取り出しているのもおかしく思えます。もしかしたらcの間違いでしょうか？　そのあたりも確認していただいて、必要なら質問のコードを修正してください。

[BoKuToTuZenU]

コードの方を修正いたしました。(参照しているvに関してはcの誤りでした。申し訳ありません。)
`assgin`のタイポについては修正させていただきました。

Answer 1

forループ内の処理をマルチスレッドで並列化するにはRayonクレートを使うのが便利です。これはforループを並列イテレーター（parallel iterator）を使った形式（例：c.par_iter().map(|item| ... ).collect()に書き換えるだけで、map()に書いた処理をマルチスレッドで並列処理してくれるものです。ここではRayonを使った方法を紹介します。

ndarrayはCargo.tomlでrayonフィーチャーを指定すると、Array3などにpar_iter()メソッドが追加され、直接Rayonの並列イテレーターを作れるようになります。ところが、ご質問のコードではこのフィーチャーは使えません。なぜならforの1回のループ内でアクセスする要素（ndarrayのスライス）が1つではなく、time - 1とtimeで指定した2つだからです。Rayonにはこういう複数の要素ずつイテレートするために、par_windows()が用意されていますが、ndarrayのrayonフィーチャーはpar_windows()をサポートしていません。

そこで、今回はcのArray3<f32>をまずVec<ArrayView2<&f32>>に変換し、そこからRayonのpar_windows()を使って2要素ずつの並列イテレーターを作ることにします。

まずは変換部分のコードです。

# Cargo.tomlにrayonを追加する

[dependencies]
ndarray = "0.14.0"
rayon = "1.5.0"

use ndarray::{arr3, s, Array, Array2, Array3};
use rayon::prelude::*;

fn cal_v(c: &Array3<f32>, dt: f32) -> Array3<f32> {
    // Rayonのpar_windowsを使うため、一番外側の次元（1次元目）でイテレートし、
    // Vec<ArrayView2<&f32>>に収集する
    let c_view2_vec: Vec<_> = c.outer_iter().collect();

    ...
}

次にpar_windows()による並列計算です。

fn cal_v(c: &Array3<f32>, dt: f32) -> Array3<f32> {
    ...

    // Rayonのpar_windowsで並列計算する。結果をVec<Array2<f32>>に収集する
    let v_array2_vec: Vec<Array2<_>> = c_view2_vec
        .par_windows(2)  // 2要素ずつイテレートする
        .map(|w| (&w[0] - &w[1]) / dt)
        .collect();

    ...
}

これによりmap内に書いた処理が複数のスレッドで並列計算されます。collect()を使って計算結果をVec<Array2<f32>>に収集します。

最後にVec<Array2<f32>>からArray3<f32>に変換します。が、私はndarrayを使うのが初めてなことがあり、簡単に書く方法がわかりませんでした。しかたがないのでご質問のコードを真似てforとassignで変換しています。

fn cal_v(c: &Array3<f32>, dt: f32) -> Array3<f32> {
    ...

    // Vec<Array2<f32>>をArray3<f32>に変換する。もっと簡単に書く方法があるかも
    let mut v = Array::zeros(c.raw_dim());
    for (time, now_v) in v_array2_vec.into_iter().enumerate() {
        v.slice_mut(s![time + 1, .., ..,]).assign(&now_v);
    }
    v
}

これで完成です。mainを含めた全体のプログラムは以下のようになりました。

// Cargo.toml
//
// [dependencies]
// ndarray = "0.14.0"
// rayon = "1.5.0"

use ndarray::{arr3, s, Array, Array2, Array3};
use rayon::prelude::*;

fn cal_v(c: &Array3<f32>, dt: f32) -> Array3<f32> {
    // Rayonのpar_windowsを使うため、一番外側の次元（1次元目）でイテレートし、
    // Vec<ArrayView2<&f32>>に収集する
    let c_view2_vec: Vec<_> = c.outer_iter().collect();

    // Rayonのpar_windowsで並列計算する。結果をVec<Array2<f32>>に収集する
    let v_array2_vec: Vec<Array2<_>> = c_view2_vec
        .par_windows(2)
        .map(|w| (&w[0] - &w[1]) / dt)
        .collect();

    // Vec<Array2<f32>>をArray3<f32>に変換する。もっと簡単に書く方法があるかも
    let mut v = Array::zeros(c.raw_dim());
    for (time, now_v) in v_array2_vec.into_iter().enumerate() {
        v.slice_mut(s![time + 1, .., ..,]).assign(&now_v);
    }
    v
}

fn main() {
    let c = arr3(&[
        [[1.0, -2.0, 3.0], [-1.0, 2.0, -3.0]],
        [[6.0, -5.0, 4.0], [-9.0, 8.0, -7.0]],
        [[-12.0, 11.0, -10.0], [15.0, -14.0, 13.0]],
    ]);
    let v = cal_v(&c, 0.25);
    println!("{}", v);

    let expected = arr3(&[
        [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
        [[-20.0, 12.0, -4.0], [32.0, -24.0, 16.0]],
        [[72.0, -64.0, 56.0], [-96.0, 88.0, -80.0]],
    ]);

    assert_eq!(v, expected);
}

追記（2021年4月10日）

上で紹介した方法は、collect()のたびに新しいVec<_>を作るため、ご質問の元のコードに比べて、その分が余分な処理になっています。

ndarrayのドキュメントをもう少し読んでみたところ、もっと効率の良いやり方があることに気づきました。以下の2点を組み合わせます。

• par_windows()ではなく、cのndarrayスライスを2つ作ってZipで対にする
• Zipのpar_map_collect()を使用する。これは、Rayonの並列イテレーター作成 → map → collectをまとめてやってくれる。

この方法ならコードが簡単になるうえ、Vec<_>を作る必要がないので効率よく処理できるはずです。

まず、par_map_collect()を使うためにndarrayを最新の0.15.xにします。またrayonフィーチャーを指定します。

# Cargo.tomlを修正する

[dependencies]
ndarray = { version = "0.15.1", features = ["rayon"] }

use文とcal_v()関数を修正します。

use ndarray::{arr3, s, Array3, Zip};

fn cal_v(c: &Array3<f32>, dt: f32) -> Array3<f32> {
    let c_len0 = c.shape()[0];

    // 1つ目のスライスは１次元目（time）の最初の要素から、最後の1つ前の要素まで
    let c_before = c.slice(s![..(c_len0 - 1), .., ..]);
    // 2つ目のスライスは１次元目（time）の2番目の要素から、最後の要素まで
    let c_now = c.slice(s![1.., .., ..]);

    // 2つのスライスをzipして、par_map_collectで並列処理する
    Zip::from(c_before)
        .and(c_now)
        .par_map_collect(|before, now| (before - now) / dt)
}

ただし、この方法を使うと、ご質問のコードとは異なり、計算結果vの１次元目の最初の要素（値が全て0）が作られなくなります。つまり、以下のmain()関数のexpectedのコメントアウトした行のデータがなくなります。

fn main() {
    let c = arr3(&[
        [[1.0, -2.0, 3.0], [-1.0, 2.0, -3.0]],
        [[6.0, -5.0, 4.0], [-9.0, 8.0, -7.0]],
        [[-12.0, 11.0, -10.0], [15.0, -14.0, 13.0]],
    ]);
    let v = cal_v(&c, 0.25);
    println!("{}", v);

    let expected = arr3(&[
        // [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], // この要素は作られなくなる
        [[-20.0, 12.0, -4.0], [32.0, -24.0, 16.0]],
        [[72.0, -64.0, 56.0], [-96.0, 88.0, -80.0]],
    ]);

    assert_eq!(v, expected);
}

Comments

[BoKuToTuZenU]

丁寧な回答をいただき誠にありがとうございます。もとの処理と異なる点などまでありがとうございます。
動作いたしました。