RustでFEM向け疎行列ソルバ HyBIT 0.8.0 を作った

久しぶりのQiita投稿です。
最近、Rustで HyBIT (Autonomous Hybrid Sparse Solver) という疎行列ソルバを開発しています。
今回0.8.0を公開したので、どういう考えで作っているのかと、現在どこまでできているのかを簡単に紹介します。

• GitHub: https://github.com/michioga/hybit
• Documentation: https://michioga.github.io/hybit/
• crates.io: https://crates.io/crates/hybit

なぜ作ったのか

FEMの大規模連立一次方程式を解いていると、

• 反復法は省メモリで高速だが、問題によっては収束しにくい
• 直接法は頑健だが、大規模化するとメモリ消費が大きい
  という問題があります。

HyBITでは、この中間を狙っています。
基本的には低コストな反復法から開始し、収束状況が悪い場合には数値的に難しい自由度を抽出します。
そして必要な局所領域だけに直接法を適用し、前処理を強化します。
概念的には、

Iterative solver
      |
      v
convergence observation
      |
      +-- good -> continue
      |
      +-- poor
            |
            v
      hard DOFs / regions
            |
            v
       local direct solve
            |
            v
      stronger preconditioner

という構成です。

HyBIT 0.8.0

0.8.0では主に以下を追加しました。

• GeneralSquare向け restarted FGMRES
• ILU(0)前処理
• analyze -> prepare -> solve-many
• ABTM G1-G7
• prepared operatorのRayon並列化
• C/C++/Fortranとの連携
• OpenMPホストとのthread数連携
  Rust 1.73以降で利用できます

use hybit::{Csr32Matrix, HybitSolver};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let a = Csr32Matrix::new(
        3,
        3,
        vec![0, 2, 5, 7],
        vec![0, 1, 0, 1, 2, 1, 2],
        vec![
             2.0, -1.0,
            -1.0,  2.0, -1.0,
                  -1.0,  2.0,
        ],
    )?;

    let b = vec![1.0, 0.0, 1.0];
    let mut x = vec![0.0; 3];

    let solver = HybitSolver::new();
    let report = solver.solve_csr32(&a, &b, &mut x)?;

    println!("x = {x:?}");
    println!("iterations = {}", report.iterations);
    println!(
        "relative residual = {:.3e}",
        report.relative_residual
    );

    Ok(())
}

入力は通常のCSR32です。
ABTMという独自形式も内部で使用していますが、利用者がABTMへ変換する必要はありません。
ABTMは疎行列形式ではなくなった
開発途中では、ABTMをCSRの代替となる高速SpMV形式として使えるのではないかと考えていました。
しかし実際にいろいろな行列で測定すると、通常のscalar SpMVでCSRを一貫して上回るものにはなりませんでした。
そこで0.8では役割を整理し、

ABTM
 |
 +-- topology
 +-- Boolean algebra
 +-- pruning
 +-- region growth
 +-- local extraction
 +-- symbolic preparation
 |
 v
CSR / Dense Block / etc.

という、topology / symbolic / preparation engineとして使う方向にしています。
これは今回の開発で得られた一つの結論です。

Rayonによる並列化

0.8ではprepared CSR operatorなどに明示的なRayon並列経路を追加しました。
一方で、小さい問題まで何でも並列化するとオーバーヘッドの方が大きくなります。
そのため、serial経路も残しています。
またC/C++/Fortranから利用する場合を考えて、OpenMPとのthread数連携も追加しました。
優先順位は、

hybit_set_num_threads()
        ↓
RAYON_NUM_THREADS
        ↓
OMP_NUM_THREADS
        ↓
Rayon default

です。

ただし、OpenMP parallel regionの各threadから同じHyBIT solveを呼ぶような使い方は、OpenMP × Rayonのnested parallelismでoversubscriptionを起こすので避けます。

実FEM問題での検証

リリース時には、実際の3次元構造FEM剛性行列を使って回帰試験を行っています。
今回の基準問題は、

自由度       : 358,065
CSR nnz      : 28,239,653
Rayon worker : 8

です。

Structural Autoでは、

iterations        : 220
verified residual : 9.378557e-9

で収束しました。

さらに同じ行列で複数RHSを解くprepared solveについても、解析結果・coarse/local factor・Krylov workspaceが再利用されることを確認しています。
性能値そのものはCPUや行列に依存するので、HyBITではrelease gateの合否には使わず、収束性・独立残差・reuseなどを基準にしています。

現在の位置づけ

HyBITはまだpre-1.0です。
現在は、

• SPD → PCG / Hybrid / structural two-level
• GeneralSquare → FGMRES
• ILU(0)
• Rayon CPU並列
  まで来ています。

一方で、

• symmetric indefinite向けMINRES
• MPI distributed solver
• production GPU backend
• out-of-core
  などはまだありません。

次はMPIホストを意識した、

partition
topology
halo
distributed execution

を進める予定です。
GPUについても考えていますが、こちらは少し後になります。

おわりに

もともとは「反復法と直接法の欠点をうまく補えないか」というところから始めた実験でした。
開発を進めるうちに、単なる新しい疎行列形式ではなく、
行列構造を解析し、問題に応じてsolver・preconditioner・executionを組み立てる基盤
という方向が見えてきました。
まだ実験的なソフトウェアですが、興味があれば試してもらえるとうれしいです。