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(¶llel, &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}