fem_structural_pcg_profile/
fem_structural_pcg_profile.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, HybitError, KrylovOutcome, LinearOperator, ParallelCsr32Operator,
11    ParallelRigidBodyTwoLevelPreconditioner, PcgWorkspace, Preconditioner, RigidBodyAggregation,
12    RigidBodyTwoLevelBlockJacobiPreconditioner, SolveStatus, 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}
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 PCG vector-kernel profile");
104    println!("Usage: fem_structural_pcg_profile --matrix K.mtx [--coords K.coords] [--rhs b.txt] [--tol 1e-8] [--max-iters 3000] [--target-coarse-dim 1536] [--aggregation graph|contiguous]");
105    println!("The profiled path uses Parallel CSR + Parallel rigid-body preconditioner and times actual PCG stages.");
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 dot_serial(a: &[f64], b: &[f64]) -> f64 {
208    debug_assert_eq!(a.len(), b.len());
209    a.iter().zip(b).map(|(x, y)| x * y).sum()
210}
211fn norm2(x: &[f64]) -> f64 {
212    dot_serial(x, x).sqrt()
213}
214fn mib(bytes: usize) -> f64 {
215    bytes as f64 / (1024.0 * 1024.0)
216}
217
218fn verified_relative_residual(
219    a: &Csr32Matrix,
220    b: &[f64],
221    x: &[f64],
222) -> Result<f64, Box<dyn Error>> {
223    let ax = a.spmv(x)?;
224    let rr = b
225        .iter()
226        .zip(&ax)
227        .map(|(&bi, &ai)| {
228            let r = bi - ai;
229            r * r
230        })
231        .sum::<f64>()
232        .sqrt();
233    let bn = norm2(b);
234    Ok(if bn == 0.0 { rr } else { rr / bn })
235}
236
237#[derive(Clone, Copy, Debug, Default)]
238struct PcgStageProfile {
239    operator_seconds: f64,
240    preconditioner_seconds: f64,
241    dot_seconds: f64,
242    norm_seconds: f64,
243    residual_init_seconds: f64,
244    update_x_r_seconds: f64,
245    update_p_seconds: f64,
246    copy_p_seconds: f64,
247    operator_calls: usize,
248    preconditioner_calls: usize,
249    dot_calls: usize,
250    norm_calls: usize,
251    update_x_r_calls: usize,
252    update_p_calls: usize,
253}
254
255impl PcgStageProfile {
256    fn accounted_seconds(&self) -> f64 {
257        self.operator_seconds
258            + self.preconditioner_seconds
259            + self.dot_seconds
260            + self.norm_seconds
261            + self.residual_init_seconds
262            + self.update_x_r_seconds
263            + self.update_p_seconds
264            + self.copy_p_seconds
265    }
266}
267
268fn timed_dot(a: &[f64], b: &[f64], profile: &mut PcgStageProfile) -> f64 {
269    let start = Instant::now();
270    let value = dot_serial(a, b);
271    profile.dot_seconds += start.elapsed().as_secs_f64();
272    profile.dot_calls += 1;
273    value
274}
275
276fn timed_norm(x: &[f64], profile: &mut PcgStageProfile) -> f64 {
277    let start = Instant::now();
278    let value = norm2(x);
279    profile.norm_seconds += start.elapsed().as_secs_f64();
280    profile.norm_calls += 1;
281    value
282}
283
284fn profiled_pcg(
285    a: &dyn LinearOperator,
286    m: &dyn Preconditioner,
287    b: &[f64],
288    x: &mut [f64],
289    options: SolverOptions,
290) -> Result<(KrylovOutcome, PcgStageProfile, f64), HybitError> {
291    options.validate()?;
292    if a.rows() != a.cols() {
293        return Err(HybitError::InvalidMatrix("PCG requires a square operator"));
294    }
295    let n = a.rows();
296    if b.len() != n {
297        return Err(HybitError::DimensionMismatch {
298            expected: n,
299            actual: b.len(),
300        });
301    }
302    if x.len() != n {
303        return Err(HybitError::DimensionMismatch {
304            expected: n,
305            actual: x.len(),
306        });
307    }
308    if m.len() != n {
309        return Err(HybitError::DimensionMismatch {
310            expected: n,
311            actual: m.len(),
312        });
313    }
314
315    let mut profile = PcgStageProfile::default();
316    let mut ax = vec![0.0; n];
317    let mut r = vec![0.0; n];
318    let mut z = vec![0.0; n];
319    let mut p = vec![0.0; n];
320    let mut ap = vec![0.0; n];
321    // Match the production prepared path: workspace allocation is setup cost,
322    // not part of the Krylov solve timer.
323    let wall_start = Instant::now();
324
325    let start = Instant::now();
326    a.apply(x, &mut ax)?;
327    profile.operator_seconds += start.elapsed().as_secs_f64();
328    profile.operator_calls += 1;
329
330    let start = Instant::now();
331    for i in 0..n {
332        r[i] = b[i] - ax[i];
333    }
334    profile.residual_init_seconds += start.elapsed().as_secs_f64();
335
336    let initial_residual = timed_norm(&r, &mut profile);
337    let b_norm = timed_norm(b, &mut profile);
338    let target = options
339        .absolute_tolerance
340        .max(options.relative_tolerance * b_norm.max(f64::MIN_POSITIVE));
341    if initial_residual <= target {
342        let wall = wall_start.elapsed().as_secs_f64();
343        return Ok((
344            KrylovOutcome {
345                status: SolveStatus::Converged,
346                iterations: 0,
347                initial_residual,
348                final_residual: initial_residual,
349            },
350            profile,
351            wall,
352        ));
353    }
354
355    let start = Instant::now();
356    m.apply(&r, &mut z)?;
357    profile.preconditioner_seconds += start.elapsed().as_secs_f64();
358    profile.preconditioner_calls += 1;
359
360    let start = Instant::now();
361    p.copy_from_slice(&z);
362    profile.copy_p_seconds += start.elapsed().as_secs_f64();
363
364    let mut rz_old = timed_dot(&r, &z, &mut profile);
365    if !rz_old.is_finite() || rz_old <= 0.0 {
366        return Err(HybitError::NumericalBreakdown(
367            "non-positive r^T M^-1 r; PCG assumptions may be violated",
368        ));
369    }
370
371    let mut final_residual = initial_residual;
372    for iter in 1..=options.max_iterations {
373        let start = Instant::now();
374        a.apply(&p, &mut ap)?;
375        profile.operator_seconds += start.elapsed().as_secs_f64();
376        profile.operator_calls += 1;
377
378        let denom = timed_dot(&p, &ap, &mut profile);
379        if !denom.is_finite() || denom <= 0.0 {
380            let wall = wall_start.elapsed().as_secs_f64();
381            return Ok((
382                KrylovOutcome {
383                    status: SolveStatus::Breakdown,
384                    iterations: iter - 1,
385                    initial_residual,
386                    final_residual,
387                },
388                profile,
389                wall,
390            ));
391        }
392
393        let alpha = rz_old / denom;
394        let start = Instant::now();
395        for i in 0..n {
396            x[i] += alpha * p[i];
397            r[i] -= alpha * ap[i];
398        }
399        profile.update_x_r_seconds += start.elapsed().as_secs_f64();
400        profile.update_x_r_calls += 1;
401
402        final_residual = timed_norm(&r, &mut profile);
403        if final_residual <= target {
404            let wall = wall_start.elapsed().as_secs_f64();
405            return Ok((
406                KrylovOutcome {
407                    status: SolveStatus::Converged,
408                    iterations: iter,
409                    initial_residual,
410                    final_residual,
411                },
412                profile,
413                wall,
414            ));
415        }
416
417        let start = Instant::now();
418        m.apply(&r, &mut z)?;
419        profile.preconditioner_seconds += start.elapsed().as_secs_f64();
420        profile.preconditioner_calls += 1;
421
422        let rz_new = timed_dot(&r, &z, &mut profile);
423        if !rz_new.is_finite() || rz_new <= 0.0 {
424            let wall = wall_start.elapsed().as_secs_f64();
425            return Ok((
426                KrylovOutcome {
427                    status: SolveStatus::Breakdown,
428                    iterations: iter,
429                    initial_residual,
430                    final_residual,
431                },
432                profile,
433                wall,
434            ));
435        }
436
437        let beta = rz_new / rz_old;
438        let start = Instant::now();
439        for i in 0..n {
440            p[i] = z[i] + beta * p[i];
441        }
442        profile.update_p_seconds += start.elapsed().as_secs_f64();
443        profile.update_p_calls += 1;
444        rz_old = rz_new;
445    }
446
447    let wall = wall_start.elapsed().as_secs_f64();
448    Ok((
449        KrylovOutcome {
450            status: SolveStatus::MaxIterations,
451            iterations: options.max_iterations,
452            initial_residual,
453            final_residual,
454        },
455        profile,
456        wall,
457    ))
458}
459
460fn print_stage(label: &str, seconds: f64, calls: usize, wall: f64) {
461    let total_ms = seconds * 1.0e3;
462    let per_call_ms = if calls == 0 {
463        0.0
464    } else {
465        total_ms / calls as f64
466    };
467    let share = if wall == 0.0 {
468        0.0
469    } else {
470        100.0 * seconds / wall
471    };
472    println!(
473        "{label:<24}: {:9.3} ms total  {:7.3} ms/call  {:5.1}%  ({} calls)",
474        total_ms, per_call_ms, share, calls
475    );
476}
477
478fn main() -> Result<(), Box<dyn Error>> {
479    let args = Args::parse()?;
480    println!(
481        "HyBIT {} PCG vector-kernel profile",
482        env!("CARGO_PKG_VERSION")
483    );
484    println!("matrix             : {}", args.matrix.display());
485    println!("coordinates        : {}", args.coordinates.display());
486
487    let load_start = Instant::now();
488    let (matrix, mm) = read_matrix_market(&args.matrix)?;
489    let matrix_load = load_start.elapsed().as_secs_f64();
490    let coord_start = Instant::now();
491    let coordinates = read_coordinates(&args.coordinates)?;
492    let coord_load = coord_start.elapsed().as_secs_f64();
493    let matrix_profile = analyze_csr32(&matrix)?;
494
495    println!(
496        "Matrix Market      : {:?}, {} input entries -> {} CSR nnz",
497        mm.symmetry, mm.input_entries, mm.csr_nnz
498    );
499    println!(
500        "dimensions         : {} x {}",
501        matrix_profile.nrows, matrix_profile.ncols
502    );
503    println!("nnz                : {}", matrix_profile.nnz);
504    println!(
505        "CSR storage        : {:.3} MiB",
506        mib(matrix.storage_bytes())
507    );
508    println!("matrix load        : {:.3} ms", matrix_load * 1.0e3);
509    println!("coordinate nodes   : {}", coordinates.len());
510    println!("coordinate load    : {:.3} ms", coord_load * 1.0e3);
511    println!("aggregation        : {:?}", args.aggregation);
512    println!("target coarse dim  : {}", args.target_coarse_dimension);
513
514    if !matrix_profile.square || !matrix_profile.full_diagonal || !matrix_profile.positive_diagonal
515    {
516        return Err(
517            "structural PCG benchmark requires a square matrix with a complete positive diagonal"
518                .into(),
519        );
520    }
521    if matrix.nrows() != coordinates.len() * 3 {
522        return Err(format!(
523            "matrix/coordinate mismatch: {} matrix rows != {} coordinate nodes * 3",
524            matrix.nrows(),
525            coordinates.len()
526        )
527        .into());
528    }
529
530    let b = if let Some(path) = args.rhs.as_deref() {
531        println!("RHS                : {}", path.display());
532        load_rhs(path, matrix.nrows())?
533    } else {
534        println!("RHS                : generated as b=A*1 (known exact solution)");
535        matrix.spmv(&vec![1.0; matrix.ncols()])?
536    };
537
538    let aggregate_nodes =
539        recommend_rigid_body_aggregate_nodes(coordinates.len(), args.target_coarse_dimension)?;
540    let setup_start = Instant::now();
541    let preconditioner = match args.aggregation {
542        RigidBodyAggregation::Graph => {
543            RigidBodyTwoLevelBlockJacobiPreconditioner::from_csr32_graph(
544                &matrix,
545                &coordinates,
546                aggregate_nodes,
547            )?
548        }
549        RigidBodyAggregation::Contiguous => RigidBodyTwoLevelBlockJacobiPreconditioner::from_csr32(
550            &matrix,
551            &coordinates,
552            aggregate_nodes,
553        )?,
554        RigidBodyAggregation::Auto => unreachable!(),
555    };
556    let parallel_preconditioner = ParallelRigidBodyTwoLevelPreconditioner::new(&preconditioner)?;
557    let parallel_operator = ParallelCsr32Operator::new(&matrix);
558    let setup_seconds = setup_start.elapsed().as_secs_f64();
559
560    println!("aggregate target   : {} nodes", aggregate_nodes);
561    println!("aggregate count    : {}", preconditioner.aggregate_count());
562    println!("coarse dimension   : {}", preconditioner.coarse_dimension());
563    println!(
564        "prec storage       : {:.3} MiB",
565        mib(preconditioner.factor_bytes())
566    );
567    println!(
568        "parallel index     : {:.3} MiB",
569        mib(parallel_preconditioner.index_storage_bytes())
570    );
571    println!("setup              : {:.3} ms", setup_seconds * 1.0e3);
572    println!(
573        "Rayon threads      : {}",
574        parallel_preconditioner.rayon_threads()
575    );
576
577    let options = SolverOptions {
578        relative_tolerance: args.relative_tolerance,
579        absolute_tolerance: 0.0,
580        max_iterations: args.max_iterations,
581    };
582
583    let mut x_profiled = vec![0.0; matrix.ncols()];
584    let (profiled_out, stages, profiled_wall) = profiled_pcg(
585        &parallel_operator,
586        &parallel_preconditioner,
587        &b,
588        &mut x_profiled,
589        options,
590    )?;
591    let profiled_verified = verified_relative_residual(&matrix, &b, &x_profiled)?;
592
593    let mut x_control = vec![0.0; matrix.ncols()];
594    let mut workspace = PcgWorkspace::new(matrix.nrows());
595    let control_start = Instant::now();
596    let control_out = pcg_with_workspace(
597        &parallel_operator,
598        &parallel_preconditioner,
599        &b,
600        &mut x_control,
601        options,
602        &mut workspace,
603    )?;
604    let control_wall = control_start.elapsed().as_secs_f64();
605    let control_verified = verified_relative_residual(&matrix, &b, &x_control)?;
606
607    println!();
608    println!("Profiled PCG");
609    println!("status             : {:?}", profiled_out.status);
610    println!("iterations         : {}", profiled_out.iterations);
611    println!(
612        "reported residual  : {:.6e}",
613        profiled_out.final_residual / norm2(&b).max(f64::MIN_POSITIVE)
614    );
615    println!("verified residual  : {:.6e}", profiled_verified);
616    println!("profiled wall      : {:.3} ms", profiled_wall * 1.0e3);
617    println!(
618        "observed / iter    : {:.3} ms",
619        if profiled_out.iterations == 0 {
620            0.0
621        } else {
622            profiled_wall * 1.0e3 / profiled_out.iterations as f64
623        }
624    );
625    println!();
626    println!("Actual nested PCG stage timings");
627    print_stage(
628        "operator A*x",
629        stages.operator_seconds,
630        stages.operator_calls,
631        profiled_wall,
632    );
633    print_stage(
634        "preconditioner",
635        stages.preconditioner_seconds,
636        stages.preconditioner_calls,
637        profiled_wall,
638    );
639    print_stage(
640        "dot reductions",
641        stages.dot_seconds,
642        stages.dot_calls,
643        profiled_wall,
644    );
645    print_stage(
646        "norm reductions",
647        stages.norm_seconds,
648        stages.norm_calls,
649        profiled_wall,
650    );
651    print_stage(
652        "x/r fused update",
653        stages.update_x_r_seconds,
654        stages.update_x_r_calls,
655        profiled_wall,
656    );
657    print_stage(
658        "p update",
659        stages.update_p_seconds,
660        stages.update_p_calls,
661        profiled_wall,
662    );
663    print_stage(
664        "initial residual",
665        stages.residual_init_seconds,
666        1,
667        profiled_wall,
668    );
669    print_stage("initial p copy", stages.copy_p_seconds, 1, profiled_wall);
670    let accounted = stages.accounted_seconds();
671    let unaccounted = (profiled_wall - accounted).max(0.0);
672    println!(
673        "accounted total        : {:9.3} ms  {:5.1}%",
674        accounted * 1.0e3,
675        100.0 * accounted / profiled_wall.max(f64::MIN_POSITIVE)
676    );
677    println!(
678        "timer/control overhead : {:9.3} ms  {:5.1}%",
679        unaccounted * 1.0e3,
680        100.0 * unaccounted / profiled_wall.max(f64::MIN_POSITIVE)
681    );
682
683    println!();
684    println!("Uninstrumented control");
685    println!("status             : {:?}", control_out.status);
686    println!("iterations         : {}", control_out.iterations);
687    println!(
688        "reported residual  : {:.6e}",
689        control_out.final_residual / norm2(&b).max(f64::MIN_POSITIVE)
690    );
691    println!("verified residual  : {:.6e}", control_verified);
692    println!("solve time         : {:.3} ms", control_wall * 1.0e3);
693    println!(
694        "profile/control    : {:.3}x",
695        profiled_wall / control_wall.max(f64::MIN_POSITIVE)
696    );
697    println!("note               : profile uses the same serial PCG vector kernels as production; only nested timers are added");
698
699    Ok(())
700}
fem_structural_pcg_profile/fem_structural_pcg_profile.rs

fem_structural_pcg_profile/
fem_structural_pcg_profile.rs
