Skip to main content

fem_structural_prepared/
fem_structural_prepared.rs

1use std::env;
2use std::error::Error;
3use std::fs::{self, File};
4use std::io::{BufRead, BufReader};
5use std::path::{Path, PathBuf};
6use std::time::Instant;
7
8use hybit::{
9    analyze_csr32, read_matrix_market, Csr32Matrix, HybitSolver, ParallelCsr32Operator,
10    RigidBodyAggregation, SolverOptions, StructuralOptions, StructuralPcgVectorPolicy,
11    StructuralPreconditionerPolicy, StructuralSpmvPolicy,
12};
13
14#[derive(Debug)]
15struct Args {
16    matrix: PathBuf,
17    coordinates: PathBuf,
18    rhs: Option<PathBuf>,
19    relative_tolerance: f64,
20    max_iterations: usize,
21    target_coarse_dimension: usize,
22    aggregation: RigidBodyAggregation,
23    spmv_policy: StructuralSpmvPolicy,
24    preconditioner_policy: StructuralPreconditionerPolicy,
25    pcg_vector_policy: StructuralPcgVectorPolicy,
26    repeats: usize,
27}
28
29impl Args {
30    fn parse() -> Result<Self, Box<dyn Error>> {
31        let mut matrix = None;
32        let mut coordinates = None;
33        let mut rhs = None;
34        let mut relative_tolerance: f64 = 1.0e-8;
35        let mut max_iterations = 3000usize;
36        let mut target_coarse_dimension = 1536usize;
37        let mut aggregation = RigidBodyAggregation::Auto;
38        let mut spmv_policy = StructuralSpmvPolicy::Auto;
39        let mut preconditioner_policy = StructuralPreconditionerPolicy::Auto;
40        let mut pcg_vector_policy = StructuralPcgVectorPolicy::Auto;
41        let mut repeats = 2usize;
42        let mut it = env::args().skip(1);
43        while let Some(arg) = it.next() {
44            match arg.as_str() {
45                "--matrix" => matrix = Some(PathBuf::from(next_value(&mut it, "--matrix")?)),
46                "--coords" => coordinates = Some(PathBuf::from(next_value(&mut it, "--coords")?)),
47                "--rhs" => rhs = Some(PathBuf::from(next_value(&mut it, "--rhs")?)),
48                "--tol" => relative_tolerance = next_value(&mut it, "--tol")?.parse()?,
49                "--max-iters" => max_iterations = next_value(&mut it, "--max-iters")?.parse()?,
50                "--target-coarse-dim" => {
51                    target_coarse_dimension = next_value(&mut it, "--target-coarse-dim")?.parse()?
52                }
53                "--aggregation" => {
54                    aggregation = match next_value(&mut it, "--aggregation")?
55                        .to_ascii_lowercase()
56                        .as_str()
57                    {
58                        "auto" => RigidBodyAggregation::Auto,
59                        "contiguous" => RigidBodyAggregation::Contiguous,
60                        "graph" => RigidBodyAggregation::Graph,
61                        other => {
62                            return Err(format!(
63                                "unknown aggregation '{other}'; use auto, contiguous, or graph"
64                            )
65                            .into())
66                        }
67                    }
68                }
69                "--spmv" => {
70                    spmv_policy = match next_value(&mut it, "--spmv")?.to_ascii_lowercase().as_str()
71                    {
72                        "auto" => StructuralSpmvPolicy::Auto,
73                        "serial" => StructuralSpmvPolicy::Serial,
74                        "parallel" => StructuralSpmvPolicy::Parallel,
75                        other => {
76                            return Err(format!(
77                                "unknown SpMV policy '{other}'; use auto, serial, or parallel"
78                            )
79                            .into())
80                        }
81                    }
82                }
83                "--precond" => {
84                    let value = next_value(&mut it, "--precond")?;
85                    preconditioner_policy = match value.to_ascii_lowercase().as_str() {
86                        "auto" => StructuralPreconditionerPolicy::Auto,
87                        "serial" => StructuralPreconditionerPolicy::Serial,
88                        "parallel" => StructuralPreconditionerPolicy::Parallel,
89                        other => {
90                            return Err(format!(
91                                "unknown preconditioner policy '{other}'; use auto, serial, or parallel"
92                            )
93                            .into())
94                        }
95                    };
96                }
97                "--pcg-vectors" => {
98                    pcg_vector_policy = match next_value(&mut it, "--pcg-vectors")?
99                        .to_ascii_lowercase()
100                        .as_str()
101                    {
102                        "auto" => StructuralPcgVectorPolicy::Auto,
103                        "serial" => StructuralPcgVectorPolicy::Serial,
104                        "parallel" => StructuralPcgVectorPolicy::Parallel,
105                        other => {
106                            return Err(format!(
107                                "unknown PCG vector policy '{other}'; use auto, serial, or parallel"
108                            )
109                            .into())
110                        }
111                    }
112                }
113                "--repeats" => repeats = next_value(&mut it, "--repeats")?.parse()?,
114                "-h" | "--help" => {
115                    print_usage();
116                    std::process::exit(0);
117                }
118                other if !other.starts_with('-') && matrix.is_none() => {
119                    matrix = Some(PathBuf::from(other))
120                }
121                other => return Err(format!("unknown argument '{other}'").into()),
122            }
123        }
124
125        let matrix = matrix.ok_or("missing matrix path; use --matrix FILE.mtx")?;
126        let coordinates = coordinates.unwrap_or_else(|| matrix.with_extension("coords"));
127        if !relative_tolerance.is_finite() || relative_tolerance <= 0.0 {
128            return Err("--tol must be finite and > 0".into());
129        }
130        if max_iterations == 0 {
131            return Err("--max-iters must be > 0".into());
132        }
133        if target_coarse_dimension < 6 {
134            return Err("--target-coarse-dim must be >= 6".into());
135        }
136        if repeats < 2 {
137            return Err("--repeats must be >= 2 so prepared reuse is exercised".into());
138        }
139
140        Ok(Self {
141            matrix,
142            coordinates,
143            rhs,
144            relative_tolerance,
145            max_iterations,
146            target_coarse_dimension,
147            aggregation,
148            spmv_policy,
149            preconditioner_policy,
150            pcg_vector_policy,
151            repeats,
152        })
153    }
154}
155
156fn next_value<I: Iterator<Item = String>>(
157    it: &mut I,
158    flag: &str,
159) -> Result<String, Box<dyn Error>> {
160    it.next()
161        .ok_or_else(|| format!("missing value after {flag}").into())
162}
163
164fn print_usage() {
165    println!("HyBIT prepared structural solve-many FEM benchmark");
166    println!("Usage: fem_structural_prepared --matrix K.mtx [--coords K.coords] [--rhs b.txt] [--tol 1e-8] [--max-iters 3000] [--target-coarse-dim 1536] [--aggregation auto|contiguous|graph] [--spmv auto|serial|parallel] [--precond auto|serial|parallel] [--pcg-vectors auto|serial|parallel] [--repeats 2]");
167}
168
169fn read_coordinates(path: &Path) -> Result<Vec<[f64; 3]>, Box<dyn Error>> {
170    let file = File::open(path)?;
171    let reader = BufReader::new(file);
172    let mut expected = None::<usize>;
173    let mut coordinates = Vec::new();
174
175    for (line_no, line) in reader.lines().enumerate() {
176        let line = line?;
177        let text = line.trim();
178        if text.is_empty() || text.starts_with('#') {
179            continue;
180        }
181        if expected.is_none() {
182            expected = Some(text.parse::<usize>().map_err(|e| {
183                format!(
184                    "{}:{}: invalid coordinate count: {e}",
185                    path.display(),
186                    line_no + 1
187                )
188            })?);
189            coordinates.reserve(expected.unwrap());
190            continue;
191        }
192        let fields: Vec<&str> = text.split_whitespace().collect();
193        if fields.len() != 3 {
194            return Err(format!(
195                "{}:{}: expected three coordinates",
196                path.display(),
197                line_no + 1
198            )
199            .into());
200        }
201        let x: f64 = fields[0].parse()?;
202        let y: f64 = fields[1].parse()?;
203        let z: f64 = fields[2].parse()?;
204        if !x.is_finite() || !y.is_finite() || !z.is_finite() {
205            return Err(
206                format!("{}:{}: non-finite coordinate", path.display(), line_no + 1).into(),
207            );
208        }
209        coordinates.push([x, y, z]);
210    }
211
212    let expected =
213        expected.ok_or_else(|| format!("{}: missing coordinate count", path.display()))?;
214    if coordinates.len() != expected {
215        return Err(format!(
216            "{}: coordinate count mismatch: header says {}, read {}",
217            path.display(),
218            expected,
219            coordinates.len()
220        )
221        .into());
222    }
223    Ok(coordinates)
224}
225
226fn load_rhs(path: &Path, n: usize) -> Result<Vec<f64>, Box<dyn Error>> {
227    let text = fs::read_to_string(path)?;
228    let mut values = Vec::with_capacity(n);
229
230    for (line_index, raw_line) in text.lines().enumerate() {
231        let line = raw_line.trim_start_matches('\u{feff}').trim();
232        if line.is_empty() || line.starts_with('#') || line.starts_with('%') {
233            continue;
234        }
235        for (token_index, token) in line.split_whitespace().enumerate() {
236            let value = token.parse::<f64>().map_err(|e| {
237                format!(
238                    "{}: invalid RHS float at line {}, token {}: {:?} ({e})",
239                    path.display(),
240                    line_index + 1,
241                    token_index + 1,
242                    token
243                )
244            })?;
245            if !value.is_finite() {
246                return Err(format!(
247                    "{}: non-finite RHS value at line {}, token {}: {:?}",
248                    path.display(),
249                    line_index + 1,
250                    token_index + 1,
251                    token
252                )
253                .into());
254            }
255            values.push(value);
256        }
257    }
258
259    if values.len() != n {
260        return Err(format!(
261            "{}: RHS length mismatch: expected {n}, got {}",
262            path.display(),
263            values.len()
264        )
265        .into());
266    }
267    Ok(values)
268}
269
270fn norm2(x: &[f64]) -> f64 {
271    x.iter().map(|v| v * v).sum::<f64>().sqrt()
272}
273
274fn verified_relative_residual(
275    a: &Csr32Matrix,
276    b: &[f64],
277    x: &[f64],
278) -> Result<f64, Box<dyn Error>> {
279    let ax = a.spmv(x)?;
280    let sum = b
281        .iter()
282        .zip(ax.iter())
283        .map(|(&bi, &ai)| {
284            let r = bi - ai;
285            r * r
286        })
287        .sum::<f64>();
288    let denom = norm2(b);
289    Ok(if denom == 0.0 {
290        sum.sqrt()
291    } else {
292        sum.sqrt() / denom
293    })
294}
295
296fn relative_error_to_ones(x: &[f64]) -> f64 {
297    let diff = x
298        .iter()
299        .map(|&xi| {
300            let d = xi - 1.0;
301            d * d
302        })
303        .sum::<f64>();
304    diff.sqrt() / (x.len() as f64).sqrt().max(f64::MIN_POSITIVE)
305}
306
307fn mib(bytes: usize) -> f64 {
308    bytes as f64 / (1024.0 * 1024.0)
309}
310
311fn main() -> Result<(), Box<dyn Error>> {
312    let args = Args::parse()?;
313    println!(
314        "HyBIT {} prepared structural solve-many FEM benchmark",
315        env!("CARGO_PKG_VERSION")
316    );
317    println!("matrix             : {}", args.matrix.display());
318    println!("coordinates        : {}", args.coordinates.display());
319
320    let load_start = Instant::now();
321    let (matrix, mm) = read_matrix_market(&args.matrix)?;
322    let matrix_load_seconds = load_start.elapsed().as_secs_f64();
323    let coord_start = Instant::now();
324    let coordinates = read_coordinates(&args.coordinates)?;
325    let coord_load_seconds = coord_start.elapsed().as_secs_f64();
326    let profile = analyze_csr32(&matrix)?;
327
328    println!(
329        "Matrix Market      : {:?}, {} input entries -> {} CSR nnz",
330        mm.symmetry, mm.input_entries, mm.csr_nnz
331    );
332    println!("dimensions         : {} x {}", profile.nrows, profile.ncols);
333    println!("nnz                : {}", profile.nnz);
334    println!(
335        "CSR storage        : {:.3} MiB",
336        mib(matrix.storage_bytes())
337    );
338    println!("matrix load        : {:.3} ms", matrix_load_seconds * 1.0e3);
339    println!("coordinate nodes   : {}", coordinates.len());
340    println!("coordinate load    : {:.3} ms", coord_load_seconds * 1.0e3);
341
342    if !profile.square || !profile.full_diagonal || !profile.positive_diagonal {
343        return Err(
344            "prepared structural PCG requires a square matrix with a complete positive diagonal"
345                .into(),
346        );
347    }
348    if matrix.nrows() != coordinates.len() * 3 {
349        return Err(format!(
350            "matrix/coordinate mismatch: {} matrix rows != {} coordinate nodes * 3",
351            matrix.nrows(),
352            coordinates.len()
353        )
354        .into());
355    }
356
357    let generated_rhs = args.rhs.is_none();
358    let b = if let Some(path) = args.rhs.as_deref() {
359        println!("RHS                : {}", path.display());
360        load_rhs(path, matrix.nrows())?
361    } else {
362        println!("RHS                : generated as b=A*1 (known exact solution)");
363        let ones = vec![1.0; matrix.ncols()];
364        matrix.spmv(&ones)?
365    };
366    println!(
367        "repeats            : {} (fresh zero initial guess each solve)",
368        args.repeats
369    );
370
371    let mut solver = HybitSolver::new();
372    solver.set_options(SolverOptions {
373        relative_tolerance: args.relative_tolerance,
374        absolute_tolerance: 0.0,
375        max_iterations: args.max_iterations,
376    })?;
377    solver.set_structural_options(StructuralOptions {
378        target_coarse_dimension: args.target_coarse_dimension,
379        aggregation: args.aggregation,
380        spmv_policy: args.spmv_policy,
381        preconditioner_policy: args.preconditioner_policy,
382        pcg_vector_policy: args.pcg_vector_policy,
383    })?;
384
385    let analysis = solver.analyze_csr32(&matrix)?;
386    let mut prepared = solver.prepare_structural_csr32(&matrix, &analysis, &coordinates)?;
387
388    println!("policy             : StructuralAuto/RigidBodyTwoLevel/Prepared");
389    println!("aggregation        : {:?}", prepared.aggregation());
390    println!("SpMV policy        : {:?}", prepared.spmv_policy());
391    println!(
392        "precond policy     : {:?}",
393        prepared.structural_preconditioner_policy()
394    );
395    println!("PCG vector policy  : {:?}", prepared.pcg_vector_policy());
396    if prepared.parallel_spmv_enabled()
397        || prepared.parallel_preconditioner_enabled()
398        || prepared.parallel_pcg_vectors_enabled()
399    {
400        println!(
401            "Rayon threads      : {}",
402            ParallelCsr32Operator::new(&matrix).rayon_threads()
403        );
404    }
405    if prepared.parallel_preconditioner_enabled() {
406        println!(
407            "parallel index     : {:.3} MiB",
408            mib(prepared.parallel_preconditioner_index_bytes())
409        );
410    }
411    println!("target coarse dim  : {}", args.target_coarse_dimension);
412    println!(
413        "aggregate nodes    : {} (auto-selected)",
414        prepared.aggregate_nodes()
415    );
416    println!("aggregate count    : {}", prepared.aggregate_count());
417    println!(
418        "aggregate min/max  : {} / {} nodes",
419        prepared.min_aggregate_nodes(),
420        prepared.max_aggregate_nodes()
421    );
422    println!("modes/aggregate    : 6");
423    println!("coarse dimension   : {}", prepared.coarse_dimension());
424    println!(
425        "base factor        : {:.3} MiB",
426        mib(prepared.base_factor_bytes())
427    );
428    println!(
429        "coarse factor      : {:.3} MiB",
430        mib(prepared.coarse_factor_bytes())
431    );
432    println!(
433        "geometry storage   : {:.3} MiB",
434        mib(prepared.geometry_bytes())
435    );
436    println!(
437        "total prec storage : {:.3} MiB",
438        mib(prepared.preconditioner_bytes())
439    );
440    println!(
441        "Krylov workspace   : {:.3} MiB",
442        mib(prepared.krylov_workspace_bytes())
443    );
444    println!(
445        "analysis once      : {:.3} ms",
446        prepared.analysis_seconds() * 1.0e3
447    );
448    println!(
449        "prepare once       : {:.3} ms",
450        prepared.prepare_seconds() * 1.0e3
451    );
452
453    let mut total_solve_seconds = 0.0f64;
454    for repetition in 1..=args.repeats {
455        let mut x = vec![0.0; matrix.ncols()];
456        let wall_start = Instant::now();
457        let report = prepared.solve(&matrix, &b, &mut x)?;
458        let wall_seconds = wall_start.elapsed().as_secs_f64();
459        let verified = verified_relative_residual(&matrix, &b, &x)?;
460        total_solve_seconds += report.solve_seconds;
461
462        println!();
463        println!("Solve #{repetition}");
464        println!("status             : {:?}", report.status);
465        println!("preconditioner     : {:?}", report.preconditioner);
466        println!("reused             : {}", report.preconditioner_reused);
467        println!("sequence           : {}", report.solve_sequence);
468        println!("iterations         : {}", report.iterations);
469        println!("reported residual  : {:.6e}", report.relative_residual);
470        println!("verified residual  : {:.6e}", verified);
471        if generated_rhs {
472            println!("relative x error   : {:.6e}", relative_error_to_ones(&x));
473        }
474        println!(
475            "analysis charged   : {:.3} ms",
476            report.analysis_seconds * 1.0e3
477        );
478        println!(
479            "prepare charged    : {:.3} ms",
480            report.prepare_seconds * 1.0e3
481        );
482        println!(
483            "solve time         : {:.3} ms",
484            report.solve_seconds * 1.0e3
485        );
486        println!("call wall          : {:.3} ms", wall_seconds * 1.0e3);
487
488        if repetition == 1 && report.preconditioner_reused {
489            return Err("first prepared structural solve unexpectedly reported reuse".into());
490        }
491        if repetition > 1 && !report.preconditioner_reused {
492            return Err(
493                "subsequent prepared structural solve did not report preconditioner reuse".into(),
494            );
495        }
496        if repetition > 1 && (report.analysis_seconds != 0.0 || report.prepare_seconds != 0.0) {
497            return Err(
498                "subsequent prepared structural solve was charged reusable setup cost".into(),
499            );
500        }
501        if !verified.is_finite() {
502            return Err("non-finite independently verified residual".into());
503        }
504    }
505
506    println!();
507    println!("Prepared reuse summary");
508    println!("solve count        : {}", prepared.solve_count());
509    println!(
510        "one-time setup     : {:.3} ms",
511        (prepared.analysis_seconds() + prepared.prepare_seconds()) * 1.0e3
512    );
513    println!("sum solve time     : {:.3} ms", total_solve_seconds * 1.0e3);
514    println!(
515        "amortized total   : {:.3} ms/solve",
516        ((prepared.analysis_seconds() + prepared.prepare_seconds() + total_solve_seconds)
517            / args.repeats as f64)
518            * 1.0e3
519    );
520
521    Ok(())
522}