Rust + Rayonによる疎行列-ベクトル積の並列化

要点

• Rustにおいて、標準のVec、Nalgebra-sparse, sprsを用いた疎行列実装を用いた疎行列-ベクトル積を実装した。
• Vecの疎行列-ベクトル積をRayonを用いて並列化した。
• 計算結果を比較したところ、著者環境では$N=1,000,000$以上の行列では並列化した計算の方が早くなった。
• sprsの計算は非常に遅いが、Nalgebra-sparseの計算はVecでの実装と同程度の速度だった。

注意

• Rustにしても疎行列にしても、初心者であるので、効率的なコードでない可能性は多いにあります。
• ここでは現状、他言語(C++など)との比較は行っていません。

疎行列と格納方法

疎行列

疎行列は、要素のうちほとんどが$0$である行列で、

\begin{pmatrix}
-2&1&0&1&0&0&0&0&0\\
1&-3&1&0&1&0&0&0&0\\
0&1&-2&0&0&1&0&0&0\\
1&0&0&-3&1&0&1&0&0\\
0&1&0&1&-4&1&0&1&0\\
0&0&1&0&1&-3&0&0&1\\
0&0&0&1&0&0&-2&1&0\\
0&0&0&0&1&0&1&-3&1\\
0&0&0&0&0&1&0&1&-2\\
\end{pmatrix}

のようなものです。実際に数値計算で扱う行列はこんなものではなく、$1000000 \times 1000000$のような非常に大きな行列を扱い、$0$以外の要素の割合ももっと少ない(0.0005 %など)です。

今回はこういった行列にベクトルをかけ合わせます。

疎行列の格納方法

疎行列はほとんどの成分が0であるため、密行列のように全成分をデータとして持つとメモリ効率も計算効率も悪くなります。そのため、疎行列専用の格納方法があります。

COO形式

疎行列では0でない成分に対して$(i,j)$成分が$a_{i,j}$であるという情報が必要なので、1成分につきこの3つの情報$(i,j,a_{ij})$を持つ形式が考えられます。これがCOO形式(Coordinate形式)で、コンピュータ上のデータとしては以下のように3つの配列

let val= vec![-2.0, 1.0, 1.0, 1.0, -3.0, 1.0, 1.0, 1.0, -2.0, 1.0, 1.0, -3.0, 1.0, 1.0, 1.0, 1.0, -4.0, 1.0, 1.0, 1.0, 1.0, -3.0, 1.0, 1.0, -2.0, 1.0, 1.0, 1.0, -3.0, 1.0, 1.0, 1.0, -2.0]
let col_index=vec![0, 1, 3, 0, 1, 2, 4, 1, 2, 5, 0, 3, 4, 6, 1, 3, 4, 5, 7, 2, 4, 5, 8, 3, 6, 7, 4, 6, 7, 8, 5, 7, 8];
let row_index=vec![0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 6, 7, 7, 7, 7, 8, 8, 8];

によって疎行列を表現します。これは前述した$12 \times 12$の行列を示しています。このように表現すれば、保持するデータは33成分*3要素=99つのデータで、配列インデックスも数値も8バイトとすれば$99\times8=792$バイトになります。これは全成分をデータとして持つ場合の$12\times 12 \times 8=1152$バイトより少ないことが分かります。この程度の大きさの行列ではうまみが少ないですが、巨大な行列ではこの差が顕著になります。

CSR方式

COO形式でも行列を圧縮できますが、col_indexは同じ要素が並んでいて無駄が多いことが分かります。これは「col_indexの0番目から1行目、3番目から2行目、7番目から3行目…」のように各行の始まる番号だけを保持する形に変更でき、row_indexの要素数を減らすことができます。

let val= vec![-2.0, 1.0, 1.0, 1.0, -3.0, 1.0, 1.0, 1.0, -2.0, 1.0, 1.0, -3.0, 1.0, 1.0, 1.0, 1.0, -4.0, 1.0, 1.0, 1.0, 1.0, -3.0, 1.0, 1.0, -2.0, 1.0, 1.0, 1.0, -3.0, 1.0, 1.0, 1.0, -2.0]
let col_index=vec![0, 1, 3, 0, 1, 2, 4, 1, 2, 5, 0, 3, 4, 6, 1, 3, 4, 5, 7, 2, 4, 5, 8, 3, 6, 7, 4, 6, 7, 8, 5, 7, 8];
let row_ptr=vec![0, 3, 7, 10, 14, 19, 23, 26, 30, 33];

これがCSR方式(Compressed Row Storage方式)です。この方式ではCOO方式よりもデータ数が(32から10に)減っています。row_ptrには、ループ処理の記述を簡単にするためにcol_indexの最終成分を示す値を追加することになっています。

CSC形式(Compressed Sparse Column方式)は、CSR方式の行方向、列方向を入れ替えたものです。

本記事では、疎行列をCOO形式で作成し、CSR形式に変換して利用しています。

Rustの疎行列ライブラリ・並列化ライブラリ

Rustはクレートと呼ばれる単位で第三者の作成したライブラリを利用することができます。ここでは疎行列ライブラリであるsprsとNalgebla-sparse、および並列化ライブラリRayonを取り上げます。

Rayon

Rayonは、Rustにおけるクレート(ライブラリ)で、データ並列化を提供しています。イメージとしては、C++のOpenMPのようなものでしょうか。

Rayonはイテレータを分割してスレッドに渡し、計算をする形をとっています。そのため、for文ではなくイテレータを使用してループを記述する必要があります。以下は、RayonのReadmeにあるサンプルです。

use rayon::prelude::*;
fn sum_of_squares(input: &[i32]) -> i32 {
    input.par_iter() // <-- just change that!
         .map(|&i| i * i)
         .sum()
}

このように、iter()をpar_iter()に変えるだけで並列化が実施できるのは便利だと感じました。

Nalgebra-sparse

Nalgebraは、線形代数計算を提供するライブラリで、多くの機能が提供されているかなり巨大なライブラリです。その中で、Nalgebra-sparseは疎行列を扱うライブラリです。

形式としてはCOO形式やCSR形式をはじめとした複数の形式をサポートしているようです。例えば、CSR方式の疎行列はrow_ptrなどのデータをVec形式で渡すことで作成できます。

let a = CsrMatrix::try_from_csr_data(
        n, //colmun_size
        n, //row_size
        row_ptr,
        col_index,
        value,
    )
    .unwrap();

もちろんちゃんと演算子が定義されているので、

let c = &a_na * &b_na;

のようにして疎行列ベクトル積を計算できます。ただ、[1]で述べられているように、Nalgebla-sparseの疎行列-ベクトル積は並列化されていないようです。今回の計算においても、CPUの負荷を見る限り並列化は実施されていないように見えました。

sprs

sprsは疎行列専用の線形代数計算ライブラリです。

使用感はNalgebla-sparseとほとんど変わらず、

let a_sp = CsMat::new(
    (n, n),
    row_ptr,
    col_index,
    val,
).to_csc();

のように疎行列を作成します。なぜかCSR方式で格納した行列の計算が異常に遅かったため、今回の計算ではCSC形式に変換して計算しています。Nalgebla-sparseと同様に、疎行列-ベクトル積は並列化されていないようです。今回の計算においても、CPUの負荷を見る限り並列化は実施されていないように見えました。

Rayonによる疎行列-ベクトル積の並列化

今回記述した疎行列-ベクトル積のコードを記載します。今回はCSR方式で格納した疎行列を、構造体で

struct SparseMatrixCsr {
    val: Vec<f64>,
    row_ptr: Vec<usize>,
    col_index: Vec<usize>,
    size: usize,
}

としています。

シーケンシャルな計算

まず、一般的な(?)for文で書いてみます。

fn sp_mv_mycsr_for(mat: &SparseMatrixCsr, b: &[f64]) -> (Vec<f64>, Duration) {
    let n = mat.size;
    let t0 = Instant::now();
    let mut c = vec![0.0; n];
    #[allow(clippy::needless_range_loop)]
    for i in 0..n {
        for j in mat.row_ptr[i]..mat.row_ptr[i + 1] {
            c[i] += mat.val[j] * b[mat.col_index[j]];
        }
    }
    let t1 = t0.elapsed();
    (c, t1)
}

今回は既報[3]と異なり、スタック領域には到底入らない大きさの配列を扱うため、全てVecとしてヒープに確保しています。RustのリンターClippyが0..nの部分でもっといい書き方があるよ、と言ってくるので、警告を消しています。実際、この場合リンターが正しくて、イテレータを使用して記述した方がRustっぽく書けます。例えば、以下のようです。

fn sp_mv_mycsr_iter(mat: &SparseMatrixCsr, b: &[f64]) -> (Vec<f64>, Duration) {
    let n = mat.size;
    let t0 = Instant::now();

    #[allow(clippy::useless_conversion)]
    let c = (0..n)
        .into_iter()
        .map(|i| {
            (mat.row_ptr[i]..mat.row_ptr[i + 1])
                .map(|j| mat.val[j] * b[mat.col_index[j]])
                .sum()
        })
        .collect();
    let t1 = t0.elapsed();
    (c, t1)
}

ここでのinto_iter()は不要なので、またClippyのサジェストをオフにしています。これは、のちにこの部分をinto_par_iter()にして並列化するため、比較としてつけています。イテレータの使い方は様々な記事があるので深く触れませんが、mapがそれぞれの要素に対して行う処理を、collectがその結果を集計する処理を行っているようなイメージです。

イテレータに関する参考記事：

イテレータを使うと、他にもいろいろな書き方が考えられ、例えば結果を入れる配列cを先に作り、その各要素をイテレータで回す

fn sp_mv_mycsr_iter2(mat: &SparseMatrixCsr, b: &[f64]) -> (Vec<f64>, Duration) {
    let n = mat.size;
    let t0 = Instant::now();
    let mut c = vec![0.0; n];
    c.iter_mut().enumerate().for_each(|(i, ci)| {
        *ci = (mat.row_ptr[i]..mat.row_ptr[i + 1])
            .map(|j| mat.val[j] * b[mat.col_index[j]])
            .sum()
    });
    let t1 = t0.elapsed();
    (c, t1)
}

のような処理も考えられます。今回の比較では、各コードを公平に比較するため、配列cのメモリ確保を計測時間に入れています。ただ、メモリ確保の行を外においても時間はほぼ変わらなかったため、計算と比較してメモリ確保の時間は非常に短いと思われます。

Rayonによる並列計算

上述のようにイテレータで書いてしまえばRayonで並列化するのは簡単で、iter()をpar_iter()に変えればよいです。内側のループを並列化するとオーバーヘッドがすごいことになるので、外側のルールを並列化すれば十分でしょう。

fn sp_mv_mycsr_par_iter(mat: &SparseMatrixCsr, b: &[f64]) -> (Vec<f64>, Duration) {
    let num_thread = num_cpus::get();
    let n = mat.size;
    let t0 = Instant::now();

    let c = (0..n)
        .into_par_iter()
        .with_min_len(n / num_thread)
        .map(|i| {
            (mat.row_ptr[i]..mat.row_ptr[i + 1])
                .map(|j| mat.val[j] * b[mat.col_index[j]])
                .sum()
        })
        .collect();
    let t1 = t0.elapsed();
    (c, t1)
}

fn sp_mv_mycsr_par_iter2(mat: &SparseMatrixCsr, b: &[f64]) -> (Vec<f64>, Duration) {
    let num_thread = num_cpus::get();    
    let n = mat.size;
    let t0 = Instant::now();
    let mut c = vec![0.0; n];
    c.par_iter_mut()
        .with_min_len(n / num_thread)
        .enumerate()
        .for_each(|(i, ci)| {
        *ci = (mat.row_ptr[i]..mat.row_ptr[i + 1])
            .map(|j| mat.val[j] * b[mat.col_index[j]])
            .sum()
    });
    let t1 = t0.elapsed();
    (c, t1)
}

Rayonの並列化はOpenMPのDynamicのようにチャンクサイズをかなり小さくとり、スレッドが空いたら処理を追加するという形になっているようです。今回のような負荷が均等な場合には行数$N$をコア数で分割するstaticな並列化の方が性能が出るはずなので、チャンクサイズ(の最小値)を.with_min_len()によって変更しています。

例えば、以下のように書けばRayonにおける各スレッドのチャンクサイズを確認できるかと思います。

let min = (0..1_000_000)
    .into_par_iter()
    .fold(|| 0, |c, _| c + 1)
    .collect::<Vec<_>>();
println!("{:?}", min); 

私の環境では1～7000までの非常に多くのスレッドが生成されているようでした。

Nalgebra-sparse・sprs

Nalgebra-sparseとsprsでの計算は1行で書けるので、ほとんど行列の形式変換コードです。

fn sp_mv_sprs(mat: &SparseMatrixCsr, b: &[f64]) -> (Vec<f64>, Duration) {
    let n = mat.size;
    let mut ind = vec![0; n];
    for (i, indi) in ind.iter_mut().enumerate() {
        *indi = i
    }
    let b_sp = CsVec::new(n, ind, b.to_vec());
    let a_sp = CsMat::new(
        (n, n),
        mat.row_ptr.clone(),
        mat.col_index.clone(),
        mat.val.clone(),
    )
    .to_csc();

    let t0 = Instant::now();
    let c = &a_sp * &b_sp;
    let t1 = t0.elapsed();
    (c.data().to_vec(), t1)
}

fn sp_mv_nalgebra(mat: &SparseMatrixCsr, b: &[f64]) -> (Vec<f64>, Duration) {
    let n = mat.size;
    let a_na = CsrMatrix::try_from_csr_data(
        n,
        n,
        mat.row_ptr.clone(),
        mat.col_index.clone(),
        mat.val.clone(),
    )
    .unwrap();

    let b_na = DVector::from_column_slice(b);

    let t0 = Instant::now();
    let c = &a_na * &b_na;
    let t1 = t0.elapsed();

    let mut c2 = vec![0.0; n];
    for i in 0..c.len() {
        c2[i] = c[i];
    }
    (c2, t1)
}

どちらのコードでも疎行列作成の際にデータを消費してしまい、次の計算に移れなくなるのでcloneして作成していますが、実際の計算ではこのような処理はすごいメモリを消費するのでやめましょう。

計算条件

計算する行列・ベクトル

今回は既報[1]と同様に、2次元等間隔直交格子を用い、2次中心差分を用いた場合のポアソン方程式を模擬し、計算を行いました。例えば$N=3\times 3$の場合、行列$\boldsymbol{A}$は一番初めに出した、以下のような疎行列になります。

\begin{pmatrix}
-2&1&0&1&0&0&0&0&0\\
1&-3&1&0&1&0&0&0&0\\
0&1&-2&0&0&1&0&0&0\\
1&0&0&-3&1&0&1&0&0\\
0&1&0&1&-4&1&0&1&0\\
0&0&1&0&1&-3&0&0&1\\
0&0&0&1&0&0&-2&1&0\\
0&0&0&0&1&0&1&-3&1\\
0&0&0&0&0&1&0&1&-2\\
\end{pmatrix}

かけ合わせるベクトル$\boldsymbol{p}$は、何でもいいのですが、ここでは$i$番目の成分が$i$であるベクトルとしています。

fn build_vector(n: usize) -> Vec<f64> {
    let b = vec![0.0; n];
    b.into_iter()
        .enumerate()
        .map(|(i, _vij)| (i + 1) as f64)
        .collect()
}

行列サイズは$N=10^2$から$N=5,000^2$までを試しました¹。

計算機

検証した計算機は以下の通りです。

• CPU:AMD Ryzen9 3900X (12 Core, 3.79 GHz)
• メモリ: 32.0 GB
• OS: Ubuntsu 22.04 in WSL2-docker of Windows 11 22H2
• コンパイラ: rustc 1.68.0

結果

結果を以下に示します。結果は10回の平均値で、単位はmsです。

$N=$	$10^2$	$20^2$	$50^2$	$100^2$	$200^2$	$500^2$	$1000^2$	$2000^2$	$5000^2$
sp_mv_mycsr_for	0.002	0.004	0.016	0.056	0.24	1.841	9.06	31.021	197.767
sp_mv_mycsr_iter	0.003	0.005	0.02	0.073	0.282	2.061	8.413	33.241	223.116
sp_mv_mycsr_iter2	0.002	0.005	0.021	0.077	0.295	2.407	9.617	37.708	240.105
sp_mv_mycsr_par_iter	3.318	2.369	3.772	4.528	3.812	4.727	10.294	23.777	139.097
sp_mv_mycsr_par_iter2	2.72	2.436	3.235	4.31	3.983	4.597	11.224	27.021	146.889
sp_mv_sprs	0.059	0.071	0.152	0.488	2.106	18.788	61.214	221.033	2502.029
sp_mv_nalgebla	0.002	0.005	0.02	0.101	0.386	2.527	9.524	39.294	278.762

考察

Nalgebra-sparse・sprsとの比較

Nalgebra-sparse・sprsの計算は、グラフを見ればシーケンシャルな計算とほぼ同じ傾向であることからも、やはり並列化はなされていないと思われます。また、Nalgebra-sparseは自分で書いた計算と同程度(多少オーバーヘッドがあるかな程度)であるのに対し、sprsの計算は非常に遅い結果となりました。しかしこれでも早くなった方で、CSR方式のまま計算するとさらに10倍以上計算時間がかかります。この結果は、既報[1]の結果とかなり異なっており、私のsprsの書き方が良くなかったのか、もしくはNalgebra-sprsに改良が入ったかのどちらかなのかなと考えています。

並列化

Rayonによる最適化は、$N$が小さい時には非常に遅い結果となりました。小さいサイズで並列化のオーバーヘッドが顕著になるのはよくあることなので、ある程度は仕方がないかなと思います。(ただ、それにしてもstaticな並列化でこんなにオーバーヘッドが大きいのか？とも思いますが……)
対して、$N$が大きい場合には並列化の計算がシーケンシャルな計算より優位になりました。本環境では、大体$N<1000^2=1,000,000$ほどから場合に並列化の恩恵が出てくるという結果になりました。とはいえ、24スレッドの計算機で速度が2倍にもなっていないので、改善の余地はあるのかなと思っています。また、イテレータの書き方による差はほぼないという結果となりました。
ちなみに、$N$が大きい場合には並列化したコードでCPUがほぼ100%まで占有されていることは確認できました。

既報[1]との比較

今回の計算は、既報[1]とほぼ同じ計算を行っているつもりです。既報[1]では、計算結果をC++のspMVと比較しています。もちろん、CPUが全く異なるため比較はできないのですが、AMDのRustコードがIntelでmklを使ったC++コードより5倍速いなんてことがあるのか……？という疑問が生じあります。今回の計算がそもそも大きな間違いをしている、[1]の記事を読み間違えているなどの可能性が考えられるので、自分の環境でMKLを使用したコードを書いたりして確認したいと考えています。本記事のどの部分でもそうですが、間違っているよ等のコメントがありましたら、ご指摘いただけると大変助かります。

最後に・感想

今回は、RustにおいてRayonによる並列化を利用した疎行列-ベクトル積を記述し、シーケンシャルな計算やライブラリを用いた結果と比較しました。結果、巨大な疎行列ではある程度並列化の恩恵が出ましたが、ある程度までの大きさならNalgebra-sparseを使用するのも実用に耐えうるのかなと感じました。

今後は共役勾配法等の行列解法を書いて、ポアソン方程式を解きたいなあと思っています。疎行列ソルバーでは russell_sparseのような外部ライブラリを使ったクレートが存在するので、これらとの速度比較もできるかなと思っています。

個人的には、Rustの勉強をして、並列化のコードを書くところまでできたのはいい経験になりました。今回初めてこういった記事を書いたため、不備などありましたら申し訳ありません。何かあればコメントなどをいただけると大変有難いです。

※本記事は個人的な内容であり、著者の所属組織とは一切関係ありません。

使用したコード

以下のリポジトリに今回作成したコードを置いています。

参考文献

以下の記事を大変参考にさせていただきました。

[1] 疎行列とベクトルを掛けたい貴方に - Hishinuma_t

[1]に関連して、COO形式-CSR形式変換に関して以下のリポジトリも参考にさせていただきました。

[2] Rustで数値計算 - termoshtt

[3] 流体計算の実行速度比較: Fortran, C++, Rust, Python, Julia - shigunodo

1. これ以上のサイズではメモリが不足してしまったようでした。理論上は$25,000,000\times 5\times 8 byte=3 GB$なので、まだ余裕があるかなとも思ったのですが、やはりcloneを多用したためメモリをたくさん消費してしまったのだと思っています。 ↩