secs_engine/mimo/
acceleration.rs

1//! Optional acceleration of the ideal-FIR shifted systems (ADR 0008).
2//! The CPU owns lambda selection, support energy and all acoustic guards.
3
4use std::sync::{Arc, Mutex};
5
6use secs_compute::CancelToken;
7use secs_types::compute::{
8    ComputeBackend, ComputePreference, ComputeReason, ComputeReport, ComputeStage, ComputeStatus,
9};
10
11use super::{design::LambdaStep, solve::RidgeSweep};
12use crate::{EngineError, decay};
13
14/// A failed optional accelerator can be replaced by a fresh CPU search.
15#[derive(Debug, Clone, thiserror::Error)]
16pub enum AccelerationError {
17    #[error("cancelled")]
18    Cancelled,
19    #[error("{detail}")]
20    Failed {
21        reason: ComputeReason,
22        detail: String,
23    },
24}
25
26impl AccelerationError {
27    pub fn failed(reason: ComputeReason, detail: impl Into<String>) -> Self {
28        Self::Failed {
29            reason,
30            detail: detail.into(),
31        }
32    }
33}
34
35/// A prepared system stays on its creating worker; Send is deliberately absent.
36pub trait ShiftedSystem {
37    fn solve(&mut self, lambda: f64, cancel: &CancelToken) -> Result<Vec<f64>, AccelerationError>;
38}
39
40/// A factory can be shared by input workers. Its prepared objects cannot.
41pub trait RidgeBackend: std::fmt::Debug + Send + Sync {
42    fn factor<'a>(
43        &'a self,
44        ata: &'a [f64],
45        atb: &'a [f64],
46        cancel: &CancelToken,
47    ) -> Result<Box<dyn ShiftedSystem + 'a>, AccelerationError>;
48
49    /// Auto may keep small systems on the CPU; an explicit GPU overrides this.
50    fn minimum_dimension(&self) -> usize {
51        0
52    }
53}
54
55#[derive(Debug, Clone, Default)]
56pub struct AccelerationSelection {
57    pub status: ComputeStatus,
58    pub backend: Option<Arc<dyn RidgeBackend>>,
59}
60
61/// The unchanged sequential Householder/refinement implementation.
62#[derive(Debug, Default)]
63pub struct CpuRidgeBackend;
64
65struct CpuSystem<'a>(RidgeSweep<'a>);
66
67impl RidgeBackend for CpuRidgeBackend {
68    fn factor<'a>(
69        &'a self,
70        ata: &'a [f64],
71        atb: &'a [f64],
72        cancel: &CancelToken,
73    ) -> Result<Box<dyn ShiftedSystem + 'a>, AccelerationError> {
74        cancelled(cancel)?;
75        if atb.is_empty()
76            || atb.len().checked_mul(atb.len()) != Some(ata.len())
77            || ata.iter().chain(atb).any(|v| !v.is_finite())
78        {
79            return Err(AccelerationError::failed(
80                ComputeReason::DeviceFailure,
81                "invalid ridge system",
82            ));
83        }
84        // The Householder tridiagonalization inside (faer `tridiag_in_place`)
85        // is one call with no place to check the token, so the token is read
86        // before and after it and one factorization bounds the cancel latency.
87        // Checking per column would mean copying faer's loop, which uses
88        // uninitialised buffers (ADR 0002). docs/MIMO.md gives measured times
89        // (audit #92 TD-96).
90        let system = CpuSystem(RidgeSweep::new(ata, atb));
91        cancelled(cancel)?;
92        Ok(Box::new(system))
93    }
94}
95
96impl ShiftedSystem for CpuSystem<'_> {
97    fn solve(&mut self, lambda: f64, cancel: &CancelToken) -> Result<Vec<f64>, AccelerationError> {
98        cancelled(cancel)?;
99        let result = self
100            .0
101            .solve(lambda)
102            .map_err(|e| AccelerationError::failed(ComputeReason::DeviceFailure, e.to_string()))?;
103        cancelled(cancel)?;
104        Ok(result)
105    }
106}
107
108fn cancelled(cancel: &CancelToken) -> Result<(), AccelerationError> {
109    cancel.check().map_err(|_| AccelerationError::Cancelled)
110}
111
112pub(crate) struct RunAcceleration<'a> {
113    selection: &'a AccelerationSelection,
114    report: Mutex<RunReport>,
115}
116
117/// The report and the fallback reason its `detail` belongs to (TD-54). The
118/// detail follows its own fallback, not `report.reason`, which can also come
119/// from the selection.
120struct RunReport {
121    report: ComputeReport,
122    detail_reason: Option<ComputeReason>,
123}
124
125impl<'a> RunAcceleration<'a> {
126    pub(crate) fn new(selection: &'a AccelerationSelection) -> Self {
127        Self {
128            selection,
129            report: Mutex::new(RunReport {
130                report: ComputeReport {
131                    requested: selection.status.preference,
132                    reason: Some(if selection.status.preference == ComputePreference::Cpu {
133                        ComputeReason::UserCpu
134                    } else {
135                        selection.status.reason
136                    }),
137                    ..ComputeReport::default()
138                },
139                detail_reason: None,
140            }),
141        }
142    }
143
144    pub(crate) fn report(&self) -> ComputeReport {
145        self.report
146            .lock()
147            .unwrap_or_else(|e| e.into_inner())
148            .report
149            .clone()
150    }
151
152    /// `fallback` is the failed accelerator's reason and detail.
153    fn completed(&self, gpu: bool, fallback: Option<(ComputeReason, &str)>, small: bool) {
154        let mut state = self.report.lock().unwrap_or_else(|e| e.into_inner());
155        let RunReport {
156            report,
157            detail_reason,
158        } = &mut *state;
159        if gpu {
160            report.accelerated_systems += 1;
161            report.backend = ComputeBackend::Cuda;
162            report.stage = Some(ComputeStage::IdealFirSweep);
163            report.device = self
164                .selection
165                .status
166                .device
167                .as_ref()
168                .map(|d| d.name.clone());
169            if report.fallbacks == 0 {
170                report.reason = None;
171            }
172        } else {
173            report.cpu_systems += 1;
174            if small && report.accelerated_systems == 0 && report.fallbacks == 0 {
175                report.reason = Some(ComputeReason::SmallProblem);
176            }
177        }
178        if let Some((reason, detail)) = fallback {
179            report.fallbacks += 1;
180            report.stage = Some(ComputeStage::IdealFirSweep);
181            report.device = self
182                .selection
183                .status
184                .device
185                .as_ref()
186                .map(|d| d.name.clone());
187            // A stable severity rule makes parallel worker reporting order irrelevant.
188            if report
189                .reason
190                .is_none_or(|old| reason_priority(reason) > reason_priority(old))
191            {
192                report.reason = Some(reason);
193            }
194            // The detail of the most severe fallback; equal severities keep
195            // the smaller text, so the worker order stays irrelevant.
196            let detail: String = detail
197                .chars()
198                .take(ComputeReport::DETAIL_MAX_CHARS)
199                .collect();
200            let replace = match (*detail_reason, &report.detail) {
201                (Some(old), Some(kept)) => {
202                    let (new, old) = (reason_priority(reason), reason_priority(old));
203                    new > old || (new == old && detail < *kept)
204                }
205                _ => true,
206            };
207            if replace {
208                *detail_reason = Some(reason);
209                report.detail = Some(detail);
210            }
211        }
212    }
213}
214
215fn reason_priority(reason: ComputeReason) -> u8 {
216    match reason {
217        ComputeReason::IirReference => 13,
218        ComputeReason::DeviceFailure => 12,
219        ComputeReason::RuntimeUnavailable => 11,
220        ComputeReason::InsufficientMemory => 10,
221        ComputeReason::UnsupportedPlatform => 9,
222        ComputeReason::NoDevice => 8,
223        ComputeReason::Integrated => 7,
224        ComputeReason::Legacy => 6,
225        ComputeReason::LowPerformance => 5,
226        ComputeReason::SmallProblem => 4,
227        ComputeReason::UserCpu => 3,
228        ComputeReason::UserGpu => 2,
229        ComputeReason::Recommended => 1,
230    }
231}
232
233pub(crate) struct RidgeProblem<'a> {
234    pub ata: &'a [f64],
235    pub atb: &'a [f64],
236    pub lp: &'a [f64],
237    pub taps: usize,
238    pub normalization: f64,
239    pub target: f64,
240    pub direct_energy: f64,
241}
242
243pub(crate) struct SweepResult {
244    pub u: Vec<f64>,
245    pub ratio: f64,
246    pub log_lambda: f64,
247    pub trajectory: Vec<LambdaStep>,
248}
249
250fn search(
251    problem: &RidgeProblem<'_>,
252    system: &mut dyn ShiftedSystem,
253    cancel: &CancelToken,
254) -> Result<SweepResult, AccelerationError> {
255    let mut evaluate = |log: f64| -> Result<(Vec<f64>, f64), AccelerationError> {
256        cancelled(cancel)?;
257        let u = system.solve(problem.normalization * 10f64.powf(log), cancel)?;
258        if u.len() != problem.atb.len() || u.iter().any(|v| !v.is_finite()) {
259            return Err(AccelerationError::failed(
260                ComputeReason::DeviceFailure,
261                "invalid accelerated solution",
262            ));
263        }
264        let mut support_energy = 0.0;
265        for k in 0..problem.atb.len() / problem.taps {
266            support_energy += decay::energy(&super::design::convolve(
267                problem.lp,
268                &u[k * problem.taps..(k + 1) * problem.taps],
269            ));
270        }
271        let ratio = support_energy / problem.direct_energy;
272        if !ratio.is_finite() {
273            return Err(AccelerationError::failed(
274                ComputeReason::DeviceFailure,
275                "nonfinite support energy",
276            ));
277        }
278        Ok((u, ratio))
279    };
280    let (mut lo, mut hi) = (-8.0, 4.0);
281    let (mut u, mut ratio) = evaluate(hi)?;
282    let mut trajectory = vec![LambdaStep {
283        mid: hi,
284        ratio,
285        accepted: ratio <= problem.target,
286    }];
287    for _ in 0..18 {
288        let mid = (lo + hi) / 2.0;
289        let (trial, r) = evaluate(mid)?;
290        trajectory.push(LambdaStep {
291            mid,
292            ratio: r,
293            accepted: r <= problem.target,
294        });
295        if r > problem.target {
296            lo = mid;
297        } else {
298            hi = mid;
299            u = trial;
300            ratio = r;
301        }
302    }
303    cancelled(cancel)?;
304    Ok(SweepResult {
305        u,
306        ratio,
307        log_lambda: hi,
308        trajectory,
309    })
310}
311
312pub(crate) fn run_sweep(
313    problem: RidgeProblem<'_>,
314    cancel: &CancelToken,
315    acceleration: Option<&RunAcceleration<'_>>,
316) -> Result<SweepResult, EngineError> {
317    let mut fallback: Option<(ComputeReason, String)> = None;
318    let mut small = false;
319    if let Some(run) = acceleration
320        && run.selection.status.preference != ComputePreference::Cpu
321        && run.selection.status.selected == ComputeBackend::Cuda
322        && let Some(backend) = &run.selection.backend
323    {
324        small = run.selection.status.preference == ComputePreference::Auto
325            && problem.atb.len() < backend.minimum_dimension();
326        if !small {
327            // The prepared GPU object/lease is dropped before a CPU restart.
328            let attempt = backend
329                .factor(problem.ata, problem.atb, cancel)
330                .and_then(|mut system| search(&problem, &mut *system, cancel));
331            match attempt {
332                Ok(result) => {
333                    run.completed(true, None, false);
334                    return Ok(result);
335                }
336                Err(AccelerationError::Cancelled) => return Err(EngineError::Cancelled),
337                Err(AccelerationError::Failed { reason, detail }) => {
338                    fallback = Some((reason, detail));
339                }
340            }
341        }
342    }
343    let cpu = CpuRidgeBackend;
344    let result = cpu
345        .factor(problem.ata, problem.atb, cancel)
346        .and_then(|mut system| search(&problem, &mut *system, cancel))
347        .map_err(|error| match error {
348            AccelerationError::Cancelled => EngineError::Cancelled,
349            AccelerationError::Failed { detail, .. } => EngineError::Input(detail),
350        })?;
351    if let Some(run) = acceleration {
352        run.completed(
353            false,
354            fallback.as_ref().map(|(r, d)| (*r, d.as_str())),
355            small,
356        );
357    }
358    Ok(result)
359}
360
361#[cfg(test)]
362mod tests {
363    use super::*;
364    use secs_types::compute::GpuDeviceInfo;
365    use std::sync::atomic::{AtomicUsize, Ordering};
366
367    #[derive(Debug)]
368    struct Fake {
369        cpu: CpuRidgeBackend,
370        calls: Arc<AtomicUsize>,
371        failure: usize,
372        cancel_at_failure: bool,
373        minimum: usize,
374    }
375    struct FakeSystem<'a> {
376        inner: Box<dyn ShiftedSystem + 'a>,
377        fake: &'a Fake,
378    }
379    impl RidgeBackend for Fake {
380        fn factor<'a>(
381            &'a self,
382            ata: &'a [f64],
383            atb: &'a [f64],
384            cancel: &CancelToken,
385        ) -> Result<Box<dyn ShiftedSystem + 'a>, AccelerationError> {
386            Ok(Box::new(FakeSystem {
387                inner: self.cpu.factor(ata, atb, cancel)?,
388                fake: self,
389            }))
390        }
391        fn minimum_dimension(&self) -> usize {
392            self.minimum
393        }
394    }
395    impl ShiftedSystem for FakeSystem<'_> {
396        fn solve(
397            &mut self,
398            lambda: f64,
399            cancel: &CancelToken,
400        ) -> Result<Vec<f64>, AccelerationError> {
401            let count = self.fake.calls.fetch_add(1, Ordering::SeqCst) + 1;
402            if count == self.fake.failure {
403                if self.fake.cancel_at_failure {
404                    cancel.cancel();
405                    return Err(AccelerationError::Cancelled);
406                }
407                return Err(AccelerationError::failed(
408                    ComputeReason::DeviceFailure,
409                    "injected device loss",
410                ));
411            }
412            self.inner.solve(lambda, cancel)
413        }
414    }
415    fn problem() -> RidgeProblem<'static> {
416        RidgeProblem {
417            ata: &[4.0, 1.0, 1.0, 3.0],
418            atb: &[1.0, 2.0],
419            lp: &[1.0],
420            taps: 1,
421            normalization: 3.5,
422            target: 0.1,
423            direct_energy: 1.0,
424        }
425    }
426    fn selection(
427        preference: ComputePreference,
428        calls: Arc<AtomicUsize>,
429        failure: usize,
430        cancel_at_failure: bool,
431        minimum: usize,
432    ) -> AccelerationSelection {
433        AccelerationSelection {
434            status: ComputeStatus {
435                preference,
436                selected: ComputeBackend::Cuda,
437                gpu_available: true,
438                reason: ComputeReason::Recommended,
439                device: Some(GpuDeviceInfo {
440                    id: 0,
441                    name: "Fake CUDA".into(),
442                    compute_major: 7,
443                    compute_minor: 5,
444                    multiprocessors: 48,
445                    memory_bytes: 8 << 30,
446                    integrated: false,
447                    supported: true,
448                    recommended: true,
449                    reason: ComputeReason::Recommended,
450                }),
451                devices: Vec::new(),
452            },
453            backend: Some(Arc::new(Fake {
454                cpu: CpuRidgeBackend,
455                calls,
456                failure,
457                cancel_at_failure,
458                minimum,
459            })),
460        }
461    }
462    #[test]
463    fn middle_failure_restarts_all_nineteen_cpu_steps_and_retains_provenance() {
464        let cancel = CancelToken::new();
465        let expected = run_sweep(problem(), &cancel, None).unwrap();
466        let calls = Arc::new(AtomicUsize::new(0));
467        let selection = selection(ComputePreference::Gpu, calls.clone(), 3, false, 0);
468        let run = RunAcceleration::new(&selection);
469        let actual = run_sweep(problem(), &cancel, Some(&run)).unwrap();
470        assert_eq!(actual.trajectory.len(), 19);
471        assert_eq!(actual.trajectory, expected.trajectory);
472        assert_eq!(actual.u, expected.u);
473        assert_eq!(calls.load(Ordering::SeqCst), 3);
474        let report = run.report();
475        assert_eq!(report.cpu_systems, 1);
476        assert_eq!(report.accelerated_systems, 0);
477        assert_eq!(report.fallbacks, 1);
478        assert_eq!(report.backend, ComputeBackend::Cpu);
479        assert_eq!(report.stage, Some(ComputeStage::IdealFirSweep));
480        assert_eq!(report.device.as_deref(), Some("Fake CUDA"));
481        assert_eq!(report.reason, Some(ComputeReason::DeviceFailure));
482        assert_eq!(report.detail.as_deref(), Some("injected device loss"));
483    }
484    #[test]
485    fn cancellation_during_shift_never_falls_back() {
486        let cancel = CancelToken::new();
487        let calls = Arc::new(AtomicUsize::new(0));
488        let selection = selection(ComputePreference::Gpu, calls.clone(), 3, true, 0);
489        let run = RunAcceleration::new(&selection);
490        assert!(matches!(
491            run_sweep(problem(), &cancel, Some(&run)),
492            Err(EngineError::Cancelled)
493        ));
494        assert_eq!(calls.load(Ordering::SeqCst), 3);
495        assert_eq!(run.report().cpu_systems, 0);
496        assert_eq!(run.report().fallbacks, 0);
497    }
498    #[test]
499    fn auto_size_floor_and_explicit_gpu_have_distinct_actual_reports() {
500        for (preference, expected_calls) in
501            [(ComputePreference::Auto, 0), (ComputePreference::Gpu, 19)]
502        {
503            let calls = Arc::new(AtomicUsize::new(0));
504            let selection = selection(preference, calls.clone(), 0, false, 10);
505            let run = RunAcceleration::new(&selection);
506            run_sweep(problem(), &CancelToken::new(), Some(&run)).unwrap();
507            assert_eq!(calls.load(Ordering::SeqCst), expected_calls);
508            let report = run.report();
509            assert_eq!(report.accelerated_systems, u32::from(expected_calls > 0));
510            assert_eq!(
511                report.reason,
512                if expected_calls == 0 {
513                    Some(ComputeReason::SmallProblem)
514                } else {
515                    None
516                }
517            );
518        }
519    }
520    #[test]
521    fn an_inconsistent_provider_cannot_override_forced_cpu_or_cpu_selection() {
522        for forced in [true, false] {
523            let calls = Arc::new(AtomicUsize::new(0));
524            let mut selection = selection(
525                if forced {
526                    ComputePreference::Cpu
527                } else {
528                    ComputePreference::Auto
529                },
530                calls.clone(),
531                0,
532                false,
533                0,
534            );
535            if !forced {
536                selection.status.selected = ComputeBackend::Cpu;
537            }
538            let run = RunAcceleration::new(&selection);
539            run_sweep(problem(), &CancelToken::new(), Some(&run)).unwrap();
540            assert_eq!(calls.load(Ordering::SeqCst), 0);
541            assert_eq!(run.report().cpu_systems, 1);
542            if forced {
543                assert_eq!(run.report().reason, Some(ComputeReason::UserCpu));
544            }
545        }
546    }
547    #[test]
548    fn parallel_failure_reason_selection_is_order_independent() {
549        let selection = selection(
550            ComputePreference::Auto,
551            Arc::new(AtomicUsize::new(0)),
552            0,
553            false,
554            0,
555        );
556        let first = RunAcceleration::new(&selection);
557        let second = RunAcceleration::new(&selection);
558        let memory = (
559            ComputeReason::InsufficientMemory,
560            "cuMemAlloc: out of memory",
561        );
562        let runtime_b = (ComputeReason::RuntimeUnavailable, "nvcuda: version 12.2");
563        let runtime_a = (ComputeReason::RuntimeUnavailable, "cuInit: 999");
564        for fallback in [memory, runtime_b, runtime_a] {
565            first.completed(false, Some(fallback), false);
566        }
567        for fallback in [runtime_a, memory, runtime_b] {
568            second.completed(false, Some(fallback), false);
569        }
570        assert_eq!(first.report(), second.report());
571        assert_eq!(
572            first.report().reason,
573            Some(ComputeReason::RuntimeUnavailable)
574        );
575        assert_eq!(first.report().detail.as_deref(), Some("cuInit: 999"));
576    }
577
578    #[test]
579    fn fallback_detail_is_bounded_to_two_hundred_characters() {
580        let selection = selection(
581            ComputePreference::Auto,
582            Arc::new(AtomicUsize::new(0)),
583            0,
584            false,
585            0,
586        );
587        let run = RunAcceleration::new(&selection);
588        let long = "é".repeat(300);
589        run.completed(false, Some((ComputeReason::DeviceFailure, &long)), false);
590        let detail = run.report().detail.unwrap();
591        assert_eq!(detail.chars().count(), ComputeReport::DETAIL_MAX_CHARS);
592        assert!(long.starts_with(&detail));
593    }
594}
595
596/// Recorded-input replay surface for hardware audits; not in production builds.
597#[cfg(feature = "test-support")]
598#[derive(Debug)]
599pub struct RidgeReplayInput<'a> {
600    pub ata: &'a [f64],
601    pub atb: &'a [f64],
602    pub lp: &'a [f64],
603    pub direct: &'a [f64],
604    pub taps: usize,
605    pub level_db: f64,
606}
607
608#[cfg(feature = "test-support")]
609#[derive(Debug)]
610pub struct RidgeReplayOutput {
611    pub u: Vec<f64>,
612    pub ratio: f64,
613    pub log_lambda: f64,
614    pub trajectory: Vec<(f64, f64, bool)>,
615}
616
617/// Runs precisely the production lambda controller with the selected solver.
618#[cfg(feature = "test-support")]
619pub fn replay_ridge(
620    input: &RidgeReplayInput<'_>,
621    backend: &dyn RidgeBackend,
622    cancel: &CancelToken,
623) -> Result<RidgeReplayOutput, AccelerationError> {
624    let n = input.atb.len();
625    if input.taps == 0
626        || !n.is_multiple_of(input.taps)
627        || n == 0
628        || n.checked_mul(n) != Some(input.ata.len())
629        || !input.level_db.is_finite()
630    {
631        return Err(AccelerationError::failed(
632            ComputeReason::DeviceFailure,
633            "invalid ridge replay inputs",
634        ));
635    }
636    let diagonal: Vec<_> = (0..n).map(|k| input.ata[k * n + k]).collect();
637    let problem = RidgeProblem {
638        ata: input.ata,
639        atb: input.atb,
640        lp: input.lp,
641        taps: input.taps,
642        normalization: secs_dsp::pairwise_sum(&diagonal) / n as f64,
643        target: 10f64.powf(input.level_db / 10.0),
644        direct_energy: decay::energy(input.direct).max(1e-30),
645    };
646    let mut system = backend.factor(input.ata, input.atb, cancel)?;
647    let result = search(&problem, &mut *system, cancel)?;
648    Ok(RidgeReplayOutput {
649        u: result.u,
650        ratio: result.ratio,
651        log_lambda: result.log_lambda,
652        trajectory: result
653            .trajectory
654            .into_iter()
655            .map(|v| (v.mid, v.ratio, v.accepted))
656            .collect(),
657    })
658}
secs_engine/mimo/acceleration.rs

secs_engine/mimo/
acceleration.rs
