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