Skip to main content

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