1use serde::{Deserialize, Serialize};
2
3use crate::{
4 assembly::AssemblySummary,
5 diagnostics::{FeaDiagnostic, FeaDiagnosticSeverity},
6 operator::{apply_k, csr_stiffness, dense_stiffness},
7 solve::preconditioner::{build_spd_preconditioner, SpdPreconditionerKind},
8 solve::{
9 backend::{kind::LinearAlgebraBackendKind, linear_algebra::LinearAlgebraBackend},
10 runtime_tensor_solver::solve_linear_system_runtime_tensor,
11 },
12};
13
14#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
15pub struct LinearSolveResult {
16 pub iterations: u32,
17 pub residual_norm: f64,
18 pub converged: bool,
19 pub host_sync_count: u32,
20 pub solver_backend: String,
21 pub device_apply_k_count: u32,
22 pub device_apply_k_attempt_count: u32,
23 pub solution: Vec<f64>,
24 pub solver_method: String,
25 pub preconditioner: String,
26 pub diagnostics: Vec<FeaDiagnostic>,
27}
28
29pub fn solve_linear_system(
30 summary: &AssemblySummary,
31 preconditioner_kind: SpdPreconditionerKind,
32 backend_kind: LinearAlgebraBackendKind,
33 algebra_backend: &dyn LinearAlgebraBackend,
34) -> LinearSolveResult {
35 let mut effective_backend_kind = backend_kind;
36 let mut runtime_tensor_fallback = false;
37 if backend_kind == LinearAlgebraBackendKind::RuntimeTensor {
38 if let Some(result) = solve_linear_system_runtime_tensor(summary, preconditioner_kind) {
39 return result;
40 }
41 effective_backend_kind = LinearAlgebraBackendKind::CpuReference;
42 runtime_tensor_fallback = true;
43 }
44
45 let has_dofs = summary.dof_count > 0;
46 if !has_dofs {
47 return LinearSolveResult {
48 iterations: 0,
49 residual_norm: 0.0,
50 converged: false,
51 host_sync_count: 0,
52 solver_backend: effective_backend_kind.as_str().to_string(),
53 device_apply_k_count: 0,
54 device_apply_k_attempt_count: 0,
55 solution: Vec::new(),
56 solver_method: "matrix_free_pcg".to_string(),
57 preconditioner: preconditioner_kind.as_str().to_string(),
58 diagnostics: vec![FeaDiagnostic {
59 code: "FEA_EMPTY_SYSTEM".to_string(),
60 severity: FeaDiagnosticSeverity::Warning,
61 message: "linear solve skipped because assembled system has zero DOFs".to_string(),
62 }],
63 };
64 }
65
66 let tol = 1.0e-8;
67 let mut initial_diagnostics = Vec::new();
68 if runtime_tensor_fallback {
69 initial_diagnostics.push(FeaDiagnostic {
70 code: "FEA_RUNTIME_TENSOR_FALLBACK".to_string(),
71 severity: FeaDiagnosticSeverity::Warning,
72 message: "runtime-tensor backend unavailable for current provider; fell back to cpu_reference solver"
73 .to_string(),
74 });
75 }
76 if let Some(dense) = dense_stiffness(&summary.operator) {
77 match solve_dense_direct(summary, dense) {
78 Ok(solution) => {
79 let b = &summary.operator.rhs;
80 let residual = algebra_backend.vec_sub(b, &apply_k(&summary.operator, &solution));
81 let residual_norm = algebra_backend.dot(&residual, &residual).sqrt();
82 let b_norm = algebra_backend.dot(b, b).sqrt().max(1.0);
83 let converged = residual_norm / b_norm <= tol;
84 let mut diagnostics = initial_diagnostics;
85 diagnostics.push(FeaDiagnostic {
86 code: "FEA_SOLVER_METHOD".to_string(),
87 severity: FeaDiagnosticSeverity::Info,
88 message:
89 "solver=dense_direct factorization=gaussian_partial_pivot matrix_free=false"
90 .to_string(),
91 });
92 if !converged {
93 diagnostics.push(FeaDiagnostic {
94 code: "FEA_DENSE_DIRECT_RESIDUAL".to_string(),
95 severity: FeaDiagnosticSeverity::Warning,
96 message: format!(
97 "dense direct solve residual_norm={residual_norm} relative_residual={}",
98 residual_norm / b_norm
99 ),
100 });
101 }
102 return LinearSolveResult {
103 iterations: 1,
104 residual_norm,
105 converged,
106 host_sync_count: 0,
107 solver_backend: effective_backend_kind.as_str().to_string(),
108 device_apply_k_count: 0,
109 device_apply_k_attempt_count: 0,
110 solution,
111 solver_method: "dense_direct".to_string(),
112 preconditioner: "none".to_string(),
113 diagnostics,
114 };
115 }
116 Err(message) => {
117 initial_diagnostics.push(FeaDiagnostic {
118 code: "FEA_DENSE_DIRECT_FALLBACK".to_string(),
119 severity: FeaDiagnosticSeverity::Warning,
120 message,
121 });
122 }
123 }
124 }
125
126 let tuned_preconditioner_kind = graph_tuned_preconditioner(summary, preconditioner_kind);
127 let max_iters = graph_tuned_max_iters(summary, 64);
128 let preconditioner = build_spd_preconditioner(summary, tuned_preconditioner_kind);
129 let traversal_order = graph_solver_traversal_order(summary);
130 let b = &summary.operator.rhs;
131 let mut x = vec![0.0; summary.dof_count];
132
133 let mut r = algebra_backend.vec_sub(b, &apply_k(&summary.operator, &x));
134 let mut z = preconditioner.apply(&r);
135 let mut p = z.clone();
136 let mut rz_old = algebra_backend.dot(&r, &z);
137 let b_norm = algebra_backend.dot(b, b).sqrt().max(1.0);
138
139 let mut diagnostics = initial_diagnostics;
140 diagnostics.push(FeaDiagnostic {
141 code: "FEA_SOLVER_METHOD".to_string(),
142 severity: FeaDiagnosticSeverity::Info,
143 message: format!(
144 "solver=pcg preconditioner={} matrix_free=true",
145 preconditioner.kind().as_str()
146 ),
147 });
148 if tuned_preconditioner_kind != preconditioner_kind {
149 diagnostics.push(FeaDiagnostic {
150 code: "FEA_GRAPH_PRECONDITIONER_TUNING".to_string(),
151 severity: FeaDiagnosticSeverity::Info,
152 message: format!(
153 "requested_preconditioner={} tuned_preconditioner={}",
154 preconditioner_kind.as_str(),
155 tuned_preconditioner_kind.as_str()
156 ),
157 });
158 }
159
160 if algebra_backend.dot(&r, &r).sqrt() / b_norm <= tol {
161 return LinearSolveResult {
162 iterations: 0,
163 residual_norm: algebra_backend.dot(&r, &r).sqrt(),
164 converged: true,
165 host_sync_count: 0,
166 solver_backend: effective_backend_kind.as_str().to_string(),
167 device_apply_k_count: 0,
168 device_apply_k_attempt_count: 0,
169 solution: x,
170 solver_method: "matrix_free_pcg".to_string(),
171 preconditioner: preconditioner.kind().as_str().to_string(),
172 diagnostics,
173 };
174 }
175
176 let mut iterations = 0u32;
177 let mut converged = false;
178 for _ in 0..max_iters {
179 let ap = apply_k(&summary.operator, &p);
180 let denom = algebra_backend.dot(&p, &ap);
181 if denom.abs() <= 1.0e-18 {
182 break;
183 }
184 let alpha = rz_old / denom;
185 algebra_backend.axpy(alpha, &p, &mut x);
186 algebra_backend.axpy(-alpha, &ap, &mut r);
187 let residual_norm = algebra_backend.dot(&r, &r).sqrt();
188 iterations += 1;
189 if residual_norm / b_norm <= tol {
190 converged = true;
191 break;
192 }
193
194 z = preconditioner.apply(&r);
195 let rz_new = algebra_backend.dot(&r, &z);
196 if rz_old.abs() <= 1.0e-18 {
197 break;
198 }
199 let beta = rz_new / rz_old;
200 for i in &traversal_order {
201 p[*i] = z[*i] + beta * p[*i];
202 }
203 rz_old = rz_new;
204 }
205
206 let residual_norm = algebra_backend.dot(&r, &r).sqrt();
207 if !converged {
208 diagnostics.push(FeaDiagnostic {
209 code: "FEA_CG_MAX_ITERS".to_string(),
210 severity: FeaDiagnosticSeverity::Warning,
211 message: format!(
212 "matrix-free cg reached max iterations ({max_iters}) with residual_norm={residual_norm}"
213 ),
214 });
215 }
216
217 LinearSolveResult {
218 iterations,
219 residual_norm,
220 converged,
221 host_sync_count: 0,
222 solver_backend: effective_backend_kind.as_str().to_string(),
223 device_apply_k_count: 0,
224 device_apply_k_attempt_count: 0,
225 solution: x,
226 solver_method: "matrix_free_pcg".to_string(),
227 preconditioner: preconditioner.kind().as_str().to_string(),
228 diagnostics,
229 }
230}
231
232fn solve_dense_direct(summary: &AssemblySummary, dense: &[f64]) -> Result<Vec<f64>, String> {
233 let n = summary.dof_count;
234 if dense.len() != n.saturating_mul(n) || summary.operator.rhs.len() != n {
235 return Err(
236 "dense direct solve skipped because matrix or rhs dimensions are invalid".to_string(),
237 );
238 }
239 let free_dofs = (0..n)
240 .filter(|idx| !summary.operator.constrained[*idx])
241 .collect::<Vec<_>>();
242 let mut x = vec![0.0_f64; n];
243 for (row, value) in x.iter_mut().enumerate().take(n) {
244 if summary.operator.constrained[row] {
245 *value = summary.operator.rhs[row];
246 }
247 }
248 if free_dofs.is_empty() {
249 return Ok(x);
250 }
251
252 let free_n = free_dofs.len();
253 let mut a = vec![0.0_f64; free_n * free_n];
254 let mut b = vec![0.0_f64; free_n];
255 for (local_row, &global_row) in free_dofs.iter().enumerate() {
256 b[local_row] = summary.operator.rhs[global_row];
257 for (local_col, &global_col) in free_dofs.iter().enumerate() {
258 a[local_row * free_n + local_col] = dense[global_row * n + global_col];
259 }
260 }
261
262 let mut free_solution = solve_dense_square(&a, &b)?;
263 for _ in 0..3 {
264 let residual = dense_residual(&a, &free_solution, &b);
265 let residual_norm = residual
266 .iter()
267 .map(|value| value * value)
268 .sum::<f64>()
269 .sqrt();
270 let rhs_norm = b
271 .iter()
272 .map(|value| value * value)
273 .sum::<f64>()
274 .sqrt()
275 .max(1.0);
276 if residual_norm / rhs_norm <= 1.0e-10 {
277 break;
278 }
279 let correction = solve_dense_square(&a, &residual)?;
280 for (value, delta) in free_solution.iter_mut().zip(correction) {
281 *value += delta;
282 }
283 }
284
285 for (local, &global) in free_dofs.iter().enumerate() {
286 x[global] = free_solution[local];
287 if !x[global].is_finite() {
288 return Err(format!(
289 "dense direct solve skipped because solution component {global} is non-finite"
290 ));
291 }
292 }
293 Ok(x)
294}
295
296fn solve_dense_square(a: &[f64], b: &[f64]) -> Result<Vec<f64>, String> {
297 let n = b.len();
298 if a.len() != n.saturating_mul(n) {
299 return Err(
300 "dense direct solve skipped because matrix or rhs dimensions are invalid".to_string(),
301 );
302 }
303 let mut a = a.to_vec();
304 let mut b = b.to_vec();
305 let row_scales = (0..n)
306 .map(|row| {
307 (0..n)
308 .map(|col| a[row * n + col].abs())
309 .fold(0.0_f64, f64::max)
310 })
311 .collect::<Vec<_>>();
312
313 for pivot in 0..n {
314 let mut pivot_row = pivot;
315 let mut pivot_score = scaled_pivot_score(&a, &row_scales, n, pivot, pivot);
316 for row in (pivot + 1)..n {
317 let candidate_score = scaled_pivot_score(&a, &row_scales, n, row, pivot);
318 if candidate_score > pivot_score {
319 pivot_score = candidate_score;
320 pivot_row = row;
321 }
322 }
323 let pivot_abs = a[pivot_row * n + pivot].abs();
324 if pivot_abs <= 1.0e-18 || !pivot_abs.is_finite() || !pivot_score.is_finite() {
325 return Err(format!(
326 "dense direct solve skipped because pivot {pivot} is singular or non-finite"
327 ));
328 }
329 if pivot_row != pivot {
330 for col in 0..n {
331 a.swap(pivot * n + col, pivot_row * n + col);
332 }
333 b.swap(pivot, pivot_row);
334 }
335
336 for row in (pivot + 1)..n {
337 let factor = a[row * n + pivot] / a[pivot * n + pivot];
338 if factor == 0.0 {
339 continue;
340 }
341 a[row * n + pivot] = 0.0;
342 for col in (pivot + 1)..n {
343 a[row * n + col] -= factor * a[pivot * n + col];
344 }
345 b[row] -= factor * b[pivot];
346 }
347 }
348
349 let mut x = vec![0.0_f64; n];
350 for row in (0..n).rev() {
351 let mut rhs = b[row];
352 for col in (row + 1)..n {
353 rhs -= a[row * n + col] * x[col];
354 }
355 let diag = a[row * n + row];
356 if diag.abs() <= 1.0e-18 || !diag.is_finite() {
357 return Err(format!(
358 "dense direct solve skipped because diagonal {row} is singular or non-finite"
359 ));
360 }
361 x[row] = rhs / diag;
362 if !x[row].is_finite() {
363 return Err(format!(
364 "dense direct solve skipped because solution component {row} is non-finite"
365 ));
366 }
367 }
368 Ok(x)
369}
370
371fn scaled_pivot_score(a: &[f64], row_scales: &[f64], n: usize, row: usize, col: usize) -> f64 {
372 let scale = row_scales[row].max(1.0e-30);
373 a[row * n + col].abs() / scale
374}
375
376fn dense_residual(a: &[f64], x: &[f64], b: &[f64]) -> Vec<f64> {
377 let n = b.len();
378 let mut residual = b.to_vec();
379 for row in 0..n {
380 let ax = (0..n).map(|col| a[row * n + col] * x[col]).sum::<f64>();
381 residual[row] -= ax;
382 }
383 residual
384}
385
386fn graph_tuned_preconditioner(
387 summary: &AssemblySummary,
388 requested: SpdPreconditionerKind,
389) -> SpdPreconditionerKind {
390 if dense_stiffness(&summary.operator).is_some() || csr_stiffness(&summary.operator).is_some() {
391 return SpdPreconditionerKind::Jacobi;
392 }
393 if let Some(graph) = summary.prep_graph_assembly.as_ref() {
394 if graph.recommend_ilu0 {
395 return SpdPreconditionerKind::Ilu0;
396 }
397 }
398 requested
399}
400
401fn graph_tuned_max_iters(summary: &AssemblySummary, base_max_iters: usize) -> usize {
402 if dense_stiffness(&summary.operator).is_some() || csr_stiffness(&summary.operator).is_some() {
403 return base_max_iters
404 .max(summary.dof_count.saturating_mul(2))
405 .max(256);
406 }
407 if let Some(graph) = summary.prep_graph_assembly.as_ref() {
408 if graph.ordering_reduction_ratio > 0.15 {
409 return ((base_max_iters as f64) * 0.85).round() as usize;
410 }
411 if graph.connected_component_count > 8 {
412 return ((base_max_iters as f64) * 1.1).round() as usize;
413 }
414 }
415 base_max_iters
416}
417
418fn graph_solver_traversal_order(summary: &AssemblySummary) -> Vec<usize> {
419 let mut order = (0..summary.dof_count).collect::<Vec<_>>();
420 if let Some(graph) = summary.prep_graph_assembly.as_ref() {
421 let seed = graph.ordering_fingerprint;
422 order.sort_by_key(|idx| {
423 let mut hash = seed ^ ((*idx as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15));
424 hash ^= hash >> 33;
425 hash = hash.wrapping_mul(0xff51afd7ed558ccd);
426 hash
427 });
428 }
429 order
430}
431
432#[cfg(test)]
433mod tests {
434 use super::*;
435 use crate::{
436 assembly::{AssemblySummary, StructuralMaterialSummary},
437 operator::OperatorSystem,
438 solve::backend::cpu_reference::CpuReferenceBackend,
439 };
440
441 fn dense_summary(dof_count: usize) -> AssemblySummary {
442 AssemblySummary {
443 dof_count,
444 structural_node_count: dof_count / 3,
445 structural_translational_dof_count: dof_count,
446 structural_rotational_dof_count: 0,
447 structural_rotation_node_count: 0,
448 structural_moment_load_count: 0,
449 structural_direct_rotational_moment_load_count: 0,
450 structural_wrench_lowering: Vec::new(),
451 structural_rotational_constraint_count: 0,
452 structural_beam_element_count: 0,
453 structural_shell_element_count: 0,
454 structural_solid_element_count: dof_count / 12,
455 structural_solid_recovery: Vec::new(),
456 structural_dof_layout: Default::default(),
457 structural_beam_recovery: Vec::new(),
458 structural_shell_recovery: Vec::new(),
459 constrained_dof_count: 0,
460 load_count: 1,
461 structural_material: StructuralMaterialSummary {
462 youngs_modulus_pa: 1.0,
463 poisson_ratio: 0.3,
464 density_kg_per_m3: 1.0,
465 lame_lambda_pa: 1.0,
466 shear_modulus_pa: 1.0,
467 },
468 prep_assembly: None,
469 prep_operator_topology: None,
470 prep_region_topology: None,
471 prep_element_assembly: None,
472 prep_element_connectivity: None,
473 prep_graph_assembly: None,
474 prep_recovery_edges: Vec::new(),
475 prep_calibration: None,
476 prep_acceptance: None,
477 prep_coordinates: None,
478 thermo_mechanical: None,
479 electro_thermal: None,
480 operator: OperatorSystem {
481 dof_count,
482 constrained: vec![false; dof_count],
483 stiffness_dense: Some(vec![0.0; dof_count * dof_count]),
484 stiffness_csr: None,
485 stiffness_diag: vec![1.0; dof_count],
486 stiffness_upper: vec![0.0; dof_count.saturating_sub(1)],
487 mass_diag: vec![1.0; dof_count],
488 damping_diag: vec![0.0; dof_count],
489 rhs: vec![1.0; dof_count],
490 },
491 }
492 }
493
494 #[test]
495 fn dense_stiffness_max_iterations_scale_with_dof_count() {
496 let summary = dense_summary(144);
497
498 assert_eq!(graph_tuned_max_iters(&summary, 64), 288);
499 }
500
501 #[test]
502 fn dense_stiffness_uses_direct_solver() {
503 let mut summary = dense_summary(2);
504 summary.operator.stiffness_dense = Some(vec![4.0, 1.0, 1.0, 3.0]);
505 summary.operator.stiffness_diag = vec![4.0, 3.0];
506 summary.operator.rhs = vec![1.0, 2.0];
507 let backend = CpuReferenceBackend;
508
509 let result = solve_linear_system(
510 &summary,
511 SpdPreconditionerKind::Jacobi,
512 LinearAlgebraBackendKind::CpuReference,
513 &backend,
514 );
515
516 assert!(result.converged);
517 assert_eq!(result.solver_method, "dense_direct");
518 assert_eq!(result.preconditioner, "none");
519 assert!((result.solution[0] - 1.0 / 11.0).abs() <= 1.0e-12);
520 assert!((result.solution[1] - 7.0 / 11.0).abs() <= 1.0e-12);
521 assert!(result
522 .diagnostics
523 .iter()
524 .any(|diag| diag.code == "FEA_SOLVER_METHOD"
525 && diag.message.contains("solver=dense_direct")));
526 }
527}