Skip to main content

fem_structural_spmv/
fem_structural_spmv.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, pcg_with_workspace, read_matrix_market, recommend_rigid_body_aggregate_nodes,
10    Csr32Matrix, LinearOperator, ParallelCsr32Operator, PcgWorkspace, RigidBodyAggregation,
11    RigidBodyTwoLevelBlockJacobiPreconditioner, SolverOptions,
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    kernel_repeats: usize,
24}
25
26impl Args {
27    fn parse() -> Result<Self, Box<dyn Error>> {
28        let mut matrix = None;
29        let mut coordinates = None;
30        let mut rhs = None;
31        let mut relative_tolerance: f64 = 1.0e-8;
32        let mut max_iterations = 3000usize;
33        let mut target_coarse_dimension = 1536usize;
34        let mut aggregation = RigidBodyAggregation::Graph;
35        let mut kernel_repeats = 20usize;
36        let mut it = env::args().skip(1);
37        while let Some(arg) = it.next() {
38            match arg.as_str() {
39                "--matrix" => matrix = Some(PathBuf::from(next_value(&mut it, "--matrix")?)),
40                "--coords" => coordinates = Some(PathBuf::from(next_value(&mut it, "--coords")?)),
41                "--rhs" => rhs = Some(PathBuf::from(next_value(&mut it, "--rhs")?)),
42                "--tol" => relative_tolerance = next_value(&mut it, "--tol")?.parse()?,
43                "--max-iters" => max_iterations = next_value(&mut it, "--max-iters")?.parse()?,
44                "--target-coarse-dim" => {
45                    target_coarse_dimension = next_value(&mut it, "--target-coarse-dim")?.parse()?
46                }
47                "--kernel-repeats" => {
48                    kernel_repeats = next_value(&mut it, "--kernel-repeats")?.parse()?
49                }
50                "--aggregation" => {
51                    aggregation = match next_value(&mut it, "--aggregation")?
52                        .to_ascii_lowercase()
53                        .as_str()
54                    {
55                        "contiguous" => RigidBodyAggregation::Contiguous,
56                        "graph" => RigidBodyAggregation::Graph,
57                        other => {
58                            return Err(format!(
59                                "unknown aggregation '{other}'; use contiguous or graph"
60                            )
61                            .into())
62                        }
63                    }
64                }
65                "-h" | "--help" => {
66                    print_usage();
67                    std::process::exit(0);
68                }
69                other if !other.starts_with('-') && matrix.is_none() => {
70                    matrix = Some(PathBuf::from(other))
71                }
72                other => return Err(format!("unknown argument '{other}'").into()),
73            }
74        }
75        let matrix = matrix.ok_or("missing matrix path; use --matrix FILE.mtx")?;
76        let coordinates = coordinates.unwrap_or_else(|| matrix.with_extension("coords"));
77        if !relative_tolerance.is_finite() || relative_tolerance <= 0.0 {
78            return Err("--tol must be finite and > 0".into());
79        }
80        if max_iterations == 0 {
81            return Err("--max-iters must be > 0".into());
82        }
83        if target_coarse_dimension < 6 {
84            return Err("--target-coarse-dim must be >= 6".into());
85        }
86        if kernel_repeats == 0 {
87            return Err("--kernel-repeats must be > 0".into());
88        }
89        Ok(Self {
90            matrix,
91            coordinates,
92            rhs,
93            relative_tolerance,
94            max_iterations,
95            target_coarse_dimension,
96            aggregation,
97            kernel_repeats,
98        })
99    }
100}
101
102fn next_value<I: Iterator<Item = String>>(
103    it: &mut I,
104    flag: &str,
105) -> Result<String, Box<dyn Error>> {
106    it.next()
107        .ok_or_else(|| format!("missing value after {flag}").into())
108}
109
110fn print_usage() {
111    println!("HyBIT structural serial vs parallel CSR SpMV benchmark");
112    println!("Usage: fem_structural_spmv --matrix K.mtx [--coords K.coords] [--rhs b.txt] [--tol 1e-8] [--max-iters 3000] [--target-coarse-dim 1536] [--aggregation graph|contiguous] [--kernel-repeats 20]");
113    println!("Set RAYON_NUM_THREADS before launch to control parallel CSR thread count.");
114}
115
116fn read_coordinates(path: &Path) -> Result<Vec<[f64; 3]>, Box<dyn Error>> {
117    let file = File::open(path)?;
118    let reader = BufReader::new(file);
119    let mut expected = None::<usize>;
120    let mut coordinates = Vec::new();
121    for (line_no, line) in reader.lines().enumerate() {
122        let line = line?;
123        let text = line.trim();
124        if text.is_empty() || text.starts_with('#') {
125            continue;
126        }
127        if expected.is_none() {
128            expected = Some(text.parse::<usize>().map_err(|e| {
129                format!(
130                    "{}:{}: invalid coordinate count: {e}",
131                    path.display(),
132                    line_no + 1
133                )
134            })?);
135            coordinates.reserve(expected.unwrap());
136            continue;
137        }
138        let fields: Vec<&str> = text.split_whitespace().collect();
139        if fields.len() != 3 {
140            return Err(format!(
141                "{}:{}: expected three coordinates",
142                path.display(),
143                line_no + 1
144            )
145            .into());
146        }
147        let xyz = [
148            fields[0].parse::<f64>()?,
149            fields[1].parse::<f64>()?,
150            fields[2].parse::<f64>()?,
151        ];
152        if xyz.iter().any(|v| !v.is_finite()) {
153            return Err(
154                format!("{}:{}: non-finite coordinate", path.display(), line_no + 1).into(),
155            );
156        }
157        coordinates.push(xyz);
158    }
159    let expected =
160        expected.ok_or_else(|| format!("{}: missing coordinate count", path.display()))?;
161    if coordinates.len() != expected {
162        return Err(format!(
163            "{}: coordinate count mismatch: header says {}, read {}",
164            path.display(),
165            expected,
166            coordinates.len()
167        )
168        .into());
169    }
170    Ok(coordinates)
171}
172
173fn load_rhs(path: &Path, n: usize) -> Result<Vec<f64>, Box<dyn Error>> {
174    let text = fs::read_to_string(path)?;
175    let mut values = Vec::with_capacity(n);
176    for (line_index, raw_line) in text.lines().enumerate() {
177        let line = raw_line.trim_start_matches('\u{feff}').trim();
178        if line.is_empty() || line.starts_with('#') || line.starts_with('%') {
179            continue;
180        }
181        for (token_index, token) in line.split_whitespace().enumerate() {
182            let value = token.parse::<f64>().map_err(|e| {
183                format!(
184                    "{}: invalid RHS float at line {}, token {}: {:?} ({e})",
185                    path.display(),
186                    line_index + 1,
187                    token_index + 1,
188                    token
189                )
190            })?;
191            if !value.is_finite() {
192                return Err(format!(
193                    "{}: non-finite RHS value at line {}, token {}",
194                    path.display(),
195                    line_index + 1,
196                    token_index + 1
197                )
198                .into());
199            }
200            values.push(value);
201        }
202    }
203    if values.len() != n {
204        return Err(format!(
205            "{}: RHS length mismatch: expected {n}, got {}",
206            path.display(),
207            values.len()
208        )
209        .into());
210    }
211    Ok(values)
212}
213
214fn norm2(x: &[f64]) -> f64 {
215    x.iter().map(|v| v * v).sum::<f64>().sqrt()
216}
217fn mib(bytes: usize) -> f64 {
218    bytes as f64 / (1024.0 * 1024.0)
219}
220fn ms_per(seconds: f64, repeats: usize) -> f64 {
221    seconds * 1.0e3 / repeats as f64
222}
223
224fn verified_relative_residual(
225    a: &Csr32Matrix,
226    b: &[f64],
227    x: &[f64],
228) -> Result<f64, Box<dyn Error>> {
229    let ax = a.spmv(x)?;
230    let rr = b
231        .iter()
232        .zip(&ax)
233        .map(|(&bi, &ai)| {
234            let r = bi - ai;
235            r * r
236        })
237        .sum::<f64>()
238        .sqrt();
239    let bn = norm2(b);
240    Ok(if bn == 0.0 { rr } else { rr / bn })
241}
242
243fn main() -> Result<(), Box<dyn Error>> {
244    let args = Args::parse()?;
245    println!(
246        "HyBIT {} serial vs parallel CSR structural benchmark",
247        env!("CARGO_PKG_VERSION")
248    );
249    println!("matrix             : {}", args.matrix.display());
250    println!("coordinates        : {}", args.coordinates.display());
251
252    let load_start = Instant::now();
253    let (matrix, mm) = read_matrix_market(&args.matrix)?;
254    let matrix_load = load_start.elapsed().as_secs_f64();
255    let coord_start = Instant::now();
256    let coordinates = read_coordinates(&args.coordinates)?;
257    let coord_load = coord_start.elapsed().as_secs_f64();
258    let profile = analyze_csr32(&matrix)?;
259
260    println!(
261        "Matrix Market      : {:?}, {} input entries -> {} CSR nnz",
262        mm.symmetry, mm.input_entries, mm.csr_nnz
263    );
264    println!("dimensions         : {} x {}", profile.nrows, profile.ncols);
265    println!("nnz                : {}", profile.nnz);
266    println!(
267        "CSR storage        : {:.3} MiB",
268        mib(matrix.storage_bytes())
269    );
270    println!("matrix load        : {:.3} ms", matrix_load * 1.0e3);
271    println!("coordinate nodes   : {}", coordinates.len());
272    println!("coordinate load    : {:.3} ms", coord_load * 1.0e3);
273    println!("aggregation        : {:?}", args.aggregation);
274    println!("target coarse dim  : {}", args.target_coarse_dimension);
275    println!("kernel repeats     : {}", args.kernel_repeats);
276
277    if !profile.square || !profile.full_diagonal || !profile.positive_diagonal {
278        return Err(
279            "structural PCG benchmark requires a square matrix with a complete positive diagonal"
280                .into(),
281        );
282    }
283    if matrix.nrows() != coordinates.len() * 3 {
284        return Err(format!(
285            "matrix/coordinate mismatch: {} matrix rows != {} coordinate nodes * 3",
286            matrix.nrows(),
287            coordinates.len()
288        )
289        .into());
290    }
291
292    let b = if let Some(path) = args.rhs.as_deref() {
293        println!("RHS                : {}", path.display());
294        load_rhs(path, matrix.nrows())?
295    } else {
296        println!("RHS                : generated as b=A*1 (known exact solution)");
297        matrix.spmv(&vec![1.0; matrix.ncols()])?
298    };
299
300    let aggregate_nodes =
301        recommend_rigid_body_aggregate_nodes(coordinates.len(), args.target_coarse_dimension)?;
302    let setup_start = Instant::now();
303    let preconditioner = match args.aggregation {
304        RigidBodyAggregation::Graph => {
305            RigidBodyTwoLevelBlockJacobiPreconditioner::from_csr32_graph(
306                &matrix,
307                &coordinates,
308                aggregate_nodes,
309            )?
310        }
311        RigidBodyAggregation::Contiguous => RigidBodyTwoLevelBlockJacobiPreconditioner::from_csr32(
312            &matrix,
313            &coordinates,
314            aggregate_nodes,
315        )?,
316        RigidBodyAggregation::Auto => unreachable!(),
317    };
318    let setup_seconds = setup_start.elapsed().as_secs_f64();
319    let parallel = ParallelCsr32Operator::new(&matrix);
320
321    println!("aggregate target   : {} nodes", aggregate_nodes);
322    println!("aggregate count    : {}", preconditioner.aggregate_count());
323    println!("coarse dimension   : {}", preconditioner.coarse_dimension());
324    println!(
325        "prec storage       : {:.3} MiB",
326        mib(preconditioner.factor_bytes())
327    );
328    println!("setup              : {:.3} ms", setup_seconds * 1.0e3);
329    println!("Rayon threads      : {}", parallel.rayon_threads());
330
331    let mut ys = vec![0.0; matrix.nrows()];
332    let mut yp = vec![0.0; matrix.nrows()];
333    matrix.apply(&b, &mut ys)?;
334    parallel.apply(&b, &mut yp)?;
335    if ys != yp {
336        return Err("serial and parallel CSR kernels produced different results".into());
337    }
338
339    let serial_start = Instant::now();
340    for _ in 0..args.kernel_repeats {
341        matrix.apply(&b, &mut ys)?;
342    }
343    let serial_spmv = serial_start.elapsed().as_secs_f64();
344
345    let parallel_start = Instant::now();
346    for _ in 0..args.kernel_repeats {
347        parallel.apply(&b, &mut yp)?;
348    }
349    let parallel_spmv = parallel_start.elapsed().as_secs_f64();
350
351    println!();
352    println!("CSR SpMV microbenchmark");
353    println!(
354        "serial / call      : {:.3} ms",
355        ms_per(serial_spmv, args.kernel_repeats)
356    );
357    println!(
358        "parallel / call    : {:.3} ms",
359        ms_per(parallel_spmv, args.kernel_repeats)
360    );
361    println!(
362        "SpMV speedup       : {:.3}x",
363        serial_spmv / parallel_spmv.max(f64::MIN_POSITIVE)
364    );
365
366    let options = SolverOptions {
367        relative_tolerance: args.relative_tolerance,
368        absolute_tolerance: 0.0,
369        max_iterations: args.max_iterations,
370    };
371
372    let mut xs = vec![0.0; matrix.ncols()];
373    let mut ws = PcgWorkspace::new(matrix.nrows());
374    let serial_solve_start = Instant::now();
375    let serial_out = pcg_with_workspace(&matrix, &preconditioner, &b, &mut xs, options, &mut ws)?;
376    let serial_solve = serial_solve_start.elapsed().as_secs_f64();
377    let serial_verified = verified_relative_residual(&matrix, &b, &xs)?;
378
379    let mut xp = vec![0.0; matrix.ncols()];
380    let mut wp = PcgWorkspace::new(matrix.nrows());
381    let parallel_solve_start = Instant::now();
382    let parallel_out =
383        pcg_with_workspace(&parallel, &preconditioner, &b, &mut xp, options, &mut wp)?;
384    let parallel_solve = parallel_solve_start.elapsed().as_secs_f64();
385    let parallel_verified = verified_relative_residual(&matrix, &b, &xp)?;
386
387    println!();
388    println!("[1/2] Serial CSR PCG");
389    println!("status             : {:?}", serial_out.status);
390    println!("iterations         : {}", serial_out.iterations);
391    println!(
392        "reported residual  : {:.6e}",
393        serial_out.final_residual / norm2(&b).max(f64::MIN_POSITIVE)
394    );
395    println!("verified residual  : {:.6e}", serial_verified);
396    println!("solve time         : {:.3} ms", serial_solve * 1.0e3);
397
398    println!();
399    println!("[2/2] Parallel CSR PCG");
400    println!("status             : {:?}", parallel_out.status);
401    println!("iterations         : {}", parallel_out.iterations);
402    println!(
403        "reported residual  : {:.6e}",
404        parallel_out.final_residual / norm2(&b).max(f64::MIN_POSITIVE)
405    );
406    println!("verified residual  : {:.6e}", parallel_verified);
407    println!("solve time         : {:.3} ms", parallel_solve * 1.0e3);
408
409    println!();
410    println!("Comparison");
411    println!(
412        "iteration ratio    : {:.3} (parallel / serial)",
413        parallel_out.iterations as f64 / serial_out.iterations.max(1) as f64
414    );
415    println!(
416        "solve-time ratio   : {:.3} (parallel / serial)",
417        parallel_solve / serial_solve.max(f64::MIN_POSITIVE)
418    );
419    println!(
420        "PCG speedup        : {:.3}x",
421        serial_solve / parallel_solve.max(f64::MIN_POSITIVE)
422    );
423
424    if !serial_verified.is_finite() || !parallel_verified.is_finite() {
425        return Err("non-finite independently verified residual".into());
426    }
427    Ok(())
428}