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 ¶llel_operator,
357 ¶llel_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 ¶llel_operator,
371 ¶llel_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}