1use std::sync::Arc;
4
5use rayon::prelude::*;
6
7use crate::compute::error::ComputeError;
8use crate::compute::simd::{simd_dot_product, simd_sum};
9
10const PARALLEL_THRESHOLD: usize = 8_192;
13
14#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
16pub struct ComputeConfig {
17 #[serde(default)]
19 pub cpu_threads: usize,
20}
21
22pub trait ComputeEngine: Send + Sync {
25 fn batch_dot_product(&self, a: &[f64], b: &[f64]) -> Result<f64, ComputeError>;
27
28 fn batch_matrix_solve(
34 &self,
35 a_flat: &[f64],
36 b: &[f64],
37 n: usize,
38 ) -> Result<Vec<f64>, ComputeError>;
39
40 fn batch_exponential_smooth(
44 &self,
45 values: &[f64],
46 alpha: f64,
47 ) -> Result<Vec<f64>, ComputeError>;
48
49 fn batch_difference(&self, values: &[f64], d: usize) -> Result<Vec<f64>, ComputeError>;
51
52 fn batch_autocorrelation(
54 &self,
55 values: &[f64],
56 max_lag: usize,
57 ) -> Result<Vec<f64>, ComputeError>;
58
59 fn batch_least_squares(
65 &self,
66 a_flat: &[f64],
67 b: &[f64],
68 m: usize,
69 n: usize,
70 ) -> Result<Vec<f64>, ComputeError>;
71}
72
73pub struct CpuEngine {
82 config: ComputeConfig,
83 pool: Arc<rayon::ThreadPool>,
84}
85
86impl CpuEngine {
87 pub fn new(config: ComputeConfig) -> Result<Self, ComputeError> {
92 let num_threads = if config.cpu_threads > 0 {
93 config.cpu_threads
94 } else {
95 std::thread::available_parallelism()
96 .map(std::num::NonZero::get)
97 .unwrap_or(4)
98 };
99
100 let pool = rayon::ThreadPoolBuilder::new()
101 .num_threads(num_threads)
102 .thread_name(|i| format!("chronix-cpu-{i}"))
103 .build()
104 .map_err(|e| {
105 ComputeError::Internal(format!("failed to build rayon thread pool: {e}"))
106 })?;
107
108 metrics::gauge!("chronix_compute_thread_pool_size").set(num_threads as f64);
109
110 Ok(Self {
111 config,
112 pool: Arc::new(pool),
113 })
114 }
115
116 #[must_use]
118 pub fn thread_pool(&self) -> &rayon::ThreadPool {
119 &self.pool
120 }
121
122 #[must_use]
124 pub fn num_threads(&self) -> usize {
125 self.pool.current_num_threads()
126 }
127
128 pub fn batch_z_score(&self, values: &[f64]) -> Result<Vec<f64>, ComputeError> {
139 if values.len() < 2 {
140 return Err(ComputeError::InsufficientData {
141 min: 2,
142 got: values.len(),
143 });
144 }
145
146 let n = values.len() as f64;
147 let mean: f64 = values.iter().sum::<f64>() / n;
148 let variance: f64 = values.iter().map(|&v| (v - mean).powi(2)).sum::<f64>() / (n - 1.0);
150 let std_dev = variance.sqrt();
151
152 if std_dev == 0.0 {
153 return Ok(vec![0.0; values.len()]);
154 }
155
156 Ok(values.iter().map(|&v| (v - mean) / std_dev).collect())
157 }
158}
159
160impl std::fmt::Debug for CpuEngine {
161 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
162 f.debug_struct("CpuEngine")
163 .field("config", &self.config)
164 .field("num_threads", &self.pool.current_num_threads())
165 .finish()
166 }
167}
168
169impl Default for CpuEngine {
170 fn default() -> Self {
171 Self::new(ComputeConfig::default()).expect("default CpuEngine config should always succeed")
172 }
173}
174
175impl ComputeEngine for CpuEngine {
176 fn batch_dot_product(&self, a: &[f64], b: &[f64]) -> Result<f64, ComputeError> {
177 metrics::counter!("chronix_compute_batch_operations_total", "operation" => "dot_product", "backend" => "cpu").increment(1);
178
179 if a.len() >= PARALLEL_THRESHOLD {
181 if a.len() != b.len() {
183 return Err(ComputeError::DimensionMismatch {
184 expected: a.len(),
185 got: b.len(),
186 });
187 }
188 let chunk_size = (a.len() / self.pool.current_num_threads()).max(1024);
189 let result: Result<f64, ComputeError> = self.pool.install(|| {
190 a.par_chunks(chunk_size)
191 .zip(b.par_chunks(chunk_size))
192 .map(|(ca, cb)| simd_dot_product(ca, cb))
193 .try_reduce(|| 0.0, |acc, v| Ok(acc + v))
194 });
195 return result;
196 }
197
198 simd_dot_product(a, b)
199 }
200
201 fn batch_matrix_solve(
202 &self,
203 a_flat: &[f64],
204 b: &[f64],
205 n: usize,
206 ) -> Result<Vec<f64>, ComputeError> {
207 if a_flat.len() != n * n {
208 return Err(ComputeError::DimensionMismatch {
209 expected: n * n,
210 got: a_flat.len(),
211 });
212 }
213 if b.len() != n {
214 return Err(ComputeError::DimensionMismatch {
215 expected: n,
216 got: b.len(),
217 });
218 }
219 metrics::counter!("chronix_compute_batch_operations_total", "operation" => "matrix_solve", "backend" => "cpu").increment(1);
220 lu_solve(a_flat, b, n)
221 }
222
223 fn batch_exponential_smooth(
224 &self,
225 values: &[f64],
226 alpha: f64,
227 ) -> Result<Vec<f64>, ComputeError> {
228 if !(0.0..=1.0).contains(&alpha) {
229 return Err(ComputeError::InvalidParameter {
230 name: "alpha",
231 value: alpha.to_string(),
232 reason: "must be in [0.0, 1.0]",
233 });
234 }
235 if values.is_empty() {
236 return Ok(Vec::new());
237 }
238
239 let mut result = Vec::with_capacity(values.len());
240 result.push(values[0]);
241 let one_minus_alpha = 1.0 - alpha;
242 for i in 1..values.len() {
243 let smoothed = alpha * values[i] + one_minus_alpha * result[i - 1];
244 result.push(smoothed);
245 }
246 Ok(result)
247 }
248
249 fn batch_difference(&self, values: &[f64], d: usize) -> Result<Vec<f64>, ComputeError> {
250 if d == 0 {
251 return Ok(values.to_vec());
252 }
253 if values.len() <= d {
254 return Err(ComputeError::InsufficientData {
255 min: d + 1,
256 got: values.len(),
257 });
258 }
259
260 let mut current = values.to_vec();
261 for _ in 0..d {
262 let mut next = Vec::with_capacity(current.len() - 1);
263 for i in 1..current.len() {
264 next.push(current[i] - current[i - 1]);
265 }
266 current = next;
267 }
268 Ok(current)
269 }
270
271 fn batch_autocorrelation(
272 &self,
273 values: &[f64],
274 max_lag: usize,
275 ) -> Result<Vec<f64>, ComputeError> {
276 if values.len() < 2 {
277 return Err(ComputeError::InsufficientData {
278 min: 2,
279 got: values.len(),
280 });
281 }
282
283 metrics::counter!("chronix_compute_batch_operations_total", "operation" => "autocorrelation", "backend" => "cpu").increment(1);
284 let n = values.len();
285 let mean = simd_sum(values) / n as f64;
286
287 let var: f64 = values.iter().map(|&v| (v - mean).powi(2)).sum();
289 if var.abs() < 1e-15 {
290 let mut result = vec![0.0; max_lag + 1];
292 result[0] = 1.0;
293 return Ok(result);
294 }
295
296 let effective_max = max_lag.min(n - 1);
297
298 let acf = if n * effective_max >= PARALLEL_THRESHOLD {
301 let pool = &self.pool;
302 pool.install(|| {
303 (0..=effective_max)
304 .into_par_iter()
305 .map(|lag| {
306 let mut cov = 0.0;
307 for i in 0..(n - lag) {
308 cov += (values[i] - mean) * (values[i + lag] - mean);
309 }
310 cov / var
311 })
312 .collect::<Vec<_>>()
313 })
314 } else {
315 let mut acf = Vec::with_capacity(effective_max + 1);
316 for lag in 0..=effective_max {
317 let mut cov = 0.0;
318 for i in 0..(n - lag) {
319 cov += (values[i] - mean) * (values[i + lag] - mean);
320 }
321 acf.push(cov / var);
322 }
323 acf
324 };
325
326 if acf.len() <= max_lag {
328 let mut padded = acf;
329 padded.resize(max_lag + 1, 0.0);
330 Ok(padded)
331 } else {
332 Ok(acf)
333 }
334 }
335
336 fn batch_least_squares(
337 &self,
338 a_flat: &[f64],
339 b: &[f64],
340 m: usize,
341 n: usize,
342 ) -> Result<Vec<f64>, ComputeError> {
343 if a_flat.len() != m * n {
344 return Err(ComputeError::DimensionMismatch {
345 expected: m * n,
346 got: a_flat.len(),
347 });
348 }
349 if b.len() != m {
350 return Err(ComputeError::DimensionMismatch {
351 expected: m,
352 got: b.len(),
353 });
354 }
355 if m < n {
356 return Err(ComputeError::DimensionMismatch {
357 expected: n,
358 got: m,
359 });
360 }
361
362 qr_least_squares(a_flat, b, m, n)
363 }
364}
365
366fn qr_least_squares(
374 a_flat: &[f64],
375 b: &[f64],
376 m: usize,
377 n: usize,
378) -> Result<Vec<f64>, ComputeError> {
379 let mut a_col = vec![0.0; m * n];
382 for i in 0..m {
383 for j in 0..n {
384 a_col[j * m + i] = a_flat[i * n + j];
385 }
386 }
387 let mut rhs = b.to_vec();
388
389 for j in 0..n {
390 let col_start = j * m + j;
392 let col_end = j * m + m;
393 let sub_col = &a_col[col_start..col_end];
394 let sub_len = m - j;
395
396 let sigma: f64 = sub_col.iter().map(|&v| v * v).sum::<f64>().sqrt();
398 if sigma < 1e-15 {
399 return Err(ComputeError::SingularMatrix);
400 }
401
402 let sign = if sub_col[0] >= 0.0 { 1.0 } else { -1.0 };
404 let mut v = sub_col.to_vec();
405 v[0] += sign * sigma;
406
407 let v_dot_v: f64 = v.iter().map(|&x| x * x).sum();
408 if v_dot_v < 1e-30 {
409 continue;
410 }
411 let tau = 2.0 / v_dot_v;
412
413 for k in j..n {
415 let ck = k * m;
416 let mut dot = 0.0;
418 for i in 0..sub_len {
419 dot += v[i] * a_col[ck + j + i];
420 }
421 let factor = tau * dot;
422 for i in 0..sub_len {
423 a_col[ck + j + i] -= factor * v[i];
424 }
425 }
426
427 let mut dot_rhs = 0.0;
429 for i in 0..sub_len {
430 dot_rhs += v[i] * rhs[j + i];
431 }
432 let factor_rhs = tau * dot_rhs;
433 for i in 0..sub_len {
434 rhs[j + i] -= factor_rhs * v[i];
435 }
436 }
437
438 let mut max_diag = 0.0f64;
442 for i in 0..n {
443 max_diag = max_diag.max(a_col[i * m + i].abs());
444 }
445 let rank_tol = max_diag * 1e-12 * (m.max(n) as f64);
446
447 let mut x = vec![0.0; n];
448 for i in (0..n).rev() {
449 let mut sum = rhs[i];
450 for k in (i + 1)..n {
451 sum -= a_col[k * m + i] * x[k];
453 }
454 let diag = a_col[i * m + i];
455 if diag.abs() < rank_tol {
456 return Err(ComputeError::SingularMatrix);
457 }
458 x[i] = sum / diag;
459 }
460
461 Ok(x)
462}
463
464fn lu_solve(a_flat: &[f64], b: &[f64], n: usize) -> Result<Vec<f64>, ComputeError> {
467 let mut a = a_flat.to_vec();
469 let rhs = b.to_vec();
470 let mut pivot = (0..n).collect::<Vec<usize>>();
471
472 for col in 0..n {
474 let mut max_val = a[pivot[col] * n + col].abs();
476 let mut max_row = col;
477 for row in (col + 1)..n {
478 let v = a[pivot[row] * n + col].abs();
479 if v > max_val {
480 max_val = v;
481 max_row = row;
482 }
483 }
484 if max_val < 1e-14 {
485 return Err(ComputeError::SingularMatrix);
486 }
487 pivot.swap(col, max_row);
488
489 let piv_row = pivot[col];
490 for row in (col + 1)..n {
491 let cur_row = pivot[row];
492 let factor = a[cur_row * n + col] / a[piv_row * n + col];
493 a[cur_row * n + col] = factor; for k in (col + 1)..n {
495 a[cur_row * n + k] -= factor * a[piv_row * n + k];
496 }
497 }
498 }
499
500 let mut y = vec![0.0; n];
502 for i in 0..n {
503 let mut sum = rhs[pivot[i]];
504 for j in 0..i {
505 sum -= a[pivot[i] * n + j] * y[j];
506 }
507 y[i] = sum;
508 }
509
510 let mut x = vec![0.0; n];
512 for i in (0..n).rev() {
513 let mut sum = y[i];
514 for j in (i + 1)..n {
515 sum -= a[pivot[i] * n + j] * x[j];
516 }
517 let diag = a[pivot[i] * n + i];
518 if diag.abs() < 1e-14 {
519 return Err(ComputeError::SingularMatrix);
520 }
521 x[i] = sum / diag;
522 }
523
524 Ok(x)
525}
526
527#[cfg(test)]
528mod tests {
529 use super::*;
530
531 fn engine() -> CpuEngine {
532 CpuEngine::default()
533 }
534
535 #[test]
536 fn dot_product_correct() {
537 let e = engine();
538 let r = e
539 .batch_dot_product(&[1.0, 2.0, 3.0], &[4.0, 5.0, 6.0])
540 .unwrap();
541 assert!((r - 32.0).abs() < 1e-10);
542 }
543
544 #[test]
545 fn dot_product_dimension_mismatch() {
546 let e = engine();
547 assert!(e.batch_dot_product(&[1.0], &[1.0, 2.0]).is_err());
548 }
549
550 #[test]
551 fn matrix_solve_identity() {
552 let e = engine();
553 let a = vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0];
555 let b = vec![5.0, 3.0, 7.0];
556 let x = e.batch_matrix_solve(&a, &b, 3).unwrap();
557 for i in 0..3 {
558 assert!((x[i] - b[i]).abs() < 1e-10);
559 }
560 }
561
562 #[test]
563 fn matrix_solve_2x2() {
564 let e = engine();
565 let a = vec![2.0, 1.0, 5.0, 3.0];
567 let b = vec![4.0, 7.0];
568 let x = e.batch_matrix_solve(&a, &b, 2).unwrap();
569 assert!((x[0] - 5.0).abs() < 1e-10);
570 assert!((x[1] - (-6.0)).abs() < 1e-10);
571 }
572
573 #[test]
574 fn matrix_solve_singular() {
575 let e = engine();
576 let a = vec![1.0, 2.0, 2.0, 4.0]; let b = vec![3.0, 6.0];
578 assert!(e.batch_matrix_solve(&a, &b, 2).is_err());
579 }
580
581 #[test]
582 fn exponential_smooth_alpha_1() {
583 let e = engine();
584 let vals = vec![1.0, 2.0, 3.0, 4.0, 5.0];
586 let result = e.batch_exponential_smooth(&vals, 1.0).unwrap();
587 assert_eq!(result.len(), vals.len());
588 for (a, b) in result.iter().zip(vals.iter()) {
589 assert!((a - b).abs() < 1e-10);
590 }
591 }
592
593 #[test]
594 fn exponential_smooth_alpha_0() {
595 let e = engine();
596 let vals = vec![10.0, 20.0, 30.0];
598 let result = e.batch_exponential_smooth(&vals, 0.0).unwrap();
599 for &v in &result {
600 assert!((v - 10.0).abs() < 1e-10);
601 }
602 }
603
604 #[test]
605 fn exponential_smooth_invalid_alpha() {
606 let e = engine();
607 assert!(e.batch_exponential_smooth(&[1.0], 1.5).is_err());
608 assert!(e.batch_exponential_smooth(&[1.0], -0.1).is_err());
609 }
610
611 #[test]
612 fn difference_order_1() {
613 let e = engine();
614 let vals = vec![1.0, 3.0, 6.0, 10.0, 15.0];
615 let diff = e.batch_difference(&vals, 1).unwrap();
616 assert_eq!(diff, vec![2.0, 3.0, 4.0, 5.0]);
617 }
618
619 #[test]
620 fn difference_order_2() {
621 let e = engine();
622 let vals = vec![1.0, 3.0, 6.0, 10.0, 15.0];
623 let diff = e.batch_difference(&vals, 2).unwrap();
624 assert_eq!(diff, vec![1.0, 1.0, 1.0]);
626 }
627
628 #[test]
629 fn difference_order_0() {
630 let e = engine();
631 let vals = vec![1.0, 2.0, 3.0];
632 let diff = e.batch_difference(&vals, 0).unwrap();
633 assert_eq!(diff, vals);
634 }
635
636 #[test]
637 fn autocorrelation_lag0_is_one() {
638 let e = engine();
639 let vals = vec![1.0, 2.0, 3.0, 2.0, 1.0, 2.0, 3.0, 2.0];
640 let acf = e.batch_autocorrelation(&vals, 3).unwrap();
641 assert!((acf[0] - 1.0).abs() < 1e-10);
642 }
643
644 #[test]
645 fn autocorrelation_constant_series() {
646 let e = engine();
647 let vals = vec![5.0; 10];
648 let acf = e.batch_autocorrelation(&vals, 3).unwrap();
649 assert!((acf[0] - 1.0).abs() < 1e-10);
650 for i in 1..=3 {
651 assert!((acf[i]).abs() < 1e-10);
652 }
653 }
654
655 #[test]
656 fn least_squares_line() {
657 let e = engine();
658 let a = vec![1.0, 0.0, 1.0, 1.0, 1.0, 2.0, 1.0, 3.0];
662 let b = vec![1.0, 3.0, 5.0, 7.0];
663 let x = e.batch_least_squares(&a, &b, 4, 2).unwrap();
664 assert!((x[0] - 1.0).abs() < 1e-10); assert!((x[1] - 2.0).abs() < 1e-10); }
667
668 #[test]
669 fn least_squares_noisy() {
670 let e = engine();
671 let a = vec![1.0, 0.0, 1.0, 1.0, 1.0, 2.0, 1.0, 3.0, 1.0, 4.0];
673 let b = vec![2.1, 4.9, 8.0, 11.1, 14.0];
674 let x = e.batch_least_squares(&a, &b, 5, 2).unwrap();
675 assert!((x[0] - 2.0).abs() < 0.5); assert!((x[1] - 3.0).abs() < 0.5); }
678
679 #[test]
680 fn least_squares_ill_conditioned() {
681 let e = engine();
685 let m = 4;
686 let n = 3;
687 let mut a = Vec::with_capacity(m * n);
689 for i in 0..m {
690 for j in 0..n {
691 a.push(1.0 / (i + j + 1) as f64);
692 }
693 }
694 let x_true = [1.0, -1.0, 1.0];
696 let mut b = vec![0.0; m];
697 for i in 0..m {
698 for j in 0..n {
699 b[i] += a[i * n + j] * x_true[j];
700 }
701 }
702 let x = e.batch_least_squares(&a, &b, m, n).unwrap();
703 for j in 0..n {
704 assert!(
705 (x[j] - x_true[j]).abs() < 1e-8,
706 "x[{j}] = {} but expected {} (diff = {})",
707 x[j],
708 x_true[j],
709 (x[j] - x_true[j]).abs(),
710 );
711 }
712 }
713
714 #[test]
715 fn least_squares_underdetermined_fails() {
716 let e = engine();
717 let a = vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0];
719 let b = vec![1.0, 2.0];
720 assert!(e.batch_least_squares(&a, &b, 2, 3).is_err());
721 }
722
723 #[test]
724 fn cpu_engine_has_dedicated_thread_pool() {
725 let config = ComputeConfig { cpu_threads: 2 };
726 let e = CpuEngine::new(config).unwrap();
727 assert_eq!(e.num_threads(), 2);
728 }
729
730 #[test]
731 fn cpu_engine_default_auto_detects_threads() {
732 let e = CpuEngine::default();
733 assert!(e.num_threads() >= 1);
734 }
735
736 #[test]
737 fn parallel_dot_product_matches_serial() {
738 let n = super::PARALLEL_THRESHOLD + 1000;
740 let a: Vec<f64> = (0..n).map(|i| (i as f64) * 0.001).collect();
741 let b: Vec<f64> = (0..n).map(|i| 1.0 - (i as f64) * 0.0005).collect();
742
743 let e = engine();
744 let result = e.batch_dot_product(&a, &b).unwrap();
745
746 let expected: f64 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
748 assert!(
749 (result - expected).abs() < expected.abs() * 1e-10,
750 "parallel dot product diverged: got {result}, expected {expected}"
751 );
752 }
753
754 #[test]
755 fn parallel_autocorrelation_matches_serial() {
756 let n = 1000;
758 let max_lag = 20;
759 let config = ComputeConfig { cpu_threads: 2 };
762 let e = CpuEngine::new(config).unwrap();
763
764 let vals: Vec<f64> = (0..n).map(|i| (i as f64 * 0.1).sin()).collect();
765 let acf = e.batch_autocorrelation(&vals, max_lag).unwrap();
766
767 assert_eq!(acf.len(), max_lag + 1);
768 assert!(
769 (acf[0] - 1.0).abs() < 1e-10,
770 "lag-0 autocorrelation must be 1.0"
771 );
772 for (lag, &v) in acf.iter().enumerate() {
774 assert!(
775 (-1.0 - 1e-10..=1.0 + 1e-10).contains(&v),
776 "acf[{lag}] = {v} out of range [-1, 1]"
777 );
778 }
779 }
780
781 #[test]
782 fn cpu_engine_debug_shows_threads() {
783 let e = CpuEngine::new(ComputeConfig { cpu_threads: 3 }).unwrap();
784 let dbg = format!("{e:?}");
785 assert!(dbg.contains("num_threads: 3"), "debug output: {dbg}");
786 }
787}