1use crate::ParShape;
11use core::num::NonZeroUsize;
12use core::time::Duration;
13
14#[cfg(feature = "threadteam")]
15pub use threadteam::{PAYS, pays_at};
16
17#[cfg(not(feature = "threadteam"))]
19pub const PAYS: f64 = 3.0;
20
21#[cfg(not(feature = "threadteam"))]
23pub fn pays_at(estimate: Duration, sections: u32, dispatch_cost: Duration) -> bool {
24 if dispatch_cost == Duration::MAX {
25 return false;
26 }
27 if sections == 0 {
28 return true;
29 }
30 estimate.as_secs_f64() >= PAYS * f64::from(sections) * dispatch_cost.as_secs_f64()
31}
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
43pub enum HostClass {
44 X86_64Avx512,
46 X86_64Avx2,
48 X86_64Scalar,
50 Aarch64Apple,
52 Aarch64,
54 Generic,
56}
57
58impl HostClass {
59 pub const ALL: [HostClass; 6] = [
61 HostClass::X86_64Avx512,
62 HostClass::X86_64Avx2,
63 HostClass::X86_64Scalar,
64 HostClass::Aarch64Apple,
65 HostClass::Aarch64,
66 HostClass::Generic,
67 ];
68
69 pub fn detect() -> Self {
72 use core::sync::atomic::{AtomicU8, Ordering};
73 static CLASS: AtomicU8 = AtomicU8::new(0);
74 let c = CLASS.load(Ordering::Relaxed);
75 if c != 0 {
76 return Self::ALL[c as usize - 1];
77 }
78 let class = Self::detect_uncached();
79 let idx = Self::ALL.iter().position(|&k| k == class).unwrap();
80 CLASS.store(idx as u8 + 1, Ordering::Relaxed);
81 class
82 }
83
84 #[cfg(all(target_arch = "x86_64", feature = "std"))]
85 fn detect_uncached() -> Self {
86 if std::arch::is_x86_feature_detected!("avx512f") {
87 HostClass::X86_64Avx512
88 } else if std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("fma") {
89 HostClass::X86_64Avx2
90 } else {
91 HostClass::X86_64Scalar
92 }
93 }
94
95 #[cfg(all(target_arch = "x86_64", not(feature = "std")))]
96 fn detect_uncached() -> Self {
97 if cfg!(target_feature = "avx512f") {
98 HostClass::X86_64Avx512
99 } else if cfg!(all(target_feature = "avx2", target_feature = "fma")) {
100 HostClass::X86_64Avx2
101 } else {
102 HostClass::X86_64Scalar
103 }
104 }
105
106 #[cfg(all(target_arch = "aarch64", target_os = "macos"))]
107 fn detect_uncached() -> Self {
108 HostClass::Aarch64Apple
109 }
110
111 #[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
112 fn detect_uncached() -> Self {
113 HostClass::Aarch64
114 }
115
116 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
117 fn detect_uncached() -> Self {
118 HostClass::Generic
119 }
120
121 pub const fn name(self) -> &'static str {
123 match self {
124 HostClass::X86_64Avx512 => "x86_64-avx512",
125 HostClass::X86_64Avx2 => "x86_64-avx2",
126 HostClass::X86_64Scalar => "x86_64-scalar",
127 HostClass::Aarch64Apple => "aarch64-apple",
128 HostClass::Aarch64 => "aarch64",
129 HostClass::Generic => "generic",
130 }
131 }
132}
133
134pub mod ns_per_mac {
137 use super::rates;
138 pub fn f64() -> f64 {
140 rates().gemm_mac.real
141 }
142 pub fn c64() -> f64 {
144 rates().gemm_mac.cplx
145 }
146 pub fn f32() -> f64 {
148 rates().gemm_single.real
149 }
150 pub fn c32() -> f64 {
152 rates().gemm_single.cplx
153 }
154}
155
156pub const NS_PER_SWAP: f64 = 0.55;
162
163#[derive(Copy, Clone, Debug, PartialEq, Eq, Default)]
171pub enum ParThreshold {
172 #[default]
174 Derived,
175 Fixed(usize),
180 Never,
182}
183
184impl From<usize> for ParThreshold {
188 fn from(value: usize) -> Self {
189 match value {
190 0 => ParThreshold::Derived,
191 usize::MAX => ParThreshold::Never,
192 n => ParThreshold::Fixed(n),
193 }
194 }
195}
196
197impl From<ParThreshold> for usize {
200 fn from(value: ParThreshold) -> Self {
201 match value {
202 ParThreshold::Derived => 0,
203 ParThreshold::Fixed(n) => n.max(1),
204 ParThreshold::Never => usize::MAX,
205 }
206 }
207}
208
209impl ParThreshold {
210 #[cfg_attr(feature = "par-trace", track_caller)]
218 pub(crate) fn region_width<T: crate::ComplexField>(
219 self,
220 par: &crate::Par,
221 units: usize,
222 rate: Rate,
223 sections: u32,
224 ) -> usize {
225 match self {
226 Self::Never => 1,
227 Self::Fixed(floor) => {
228 if units >= floor {
229 par.degree()
230 } else {
231 1
232 }
233 },
234 Self::Derived => par_width::<T>(par, units, rate, sections),
235 }
236 }
237}
238
239#[derive(Copy, Clone, Debug, PartialEq)]
245pub struct Rate {
246 pub real: f64,
248 pub cplx: f64,
250}
251
252impl Rate {
253 pub fn of<T: crate::ComplexField>(self) -> f64 {
255 if const { T::IS_NATIVE_F32 } {
256 self.real / 2.0
257 } else if const { T::IS_NATIVE_C32 } {
258 self.cplx / 2.0
259 } else if const { T::IS_REAL } {
260 self.real
261 } else {
262 self.cplx
263 }
264 }
265}
266
267#[derive(Copy, Clone, Debug, PartialEq)]
289pub struct Rates {
290 pub gemm_mac: Rate,
293 pub gemm_single: Rate,
295 pub block_apply_mac: Rate,
297 pub lu_full_pivot: Rate,
300 pub bunch_kaufman: Rate,
303 pub qr_col_pivot: Rate,
306 pub bidiag: Rate,
309 pub hessenberg: Rate,
312 pub tridiag: Rate,
315 pub matvec: Rate,
318 pub bidiag_svd_dc: Rate,
325 pub tridiag_evd_dc: Rate,
329 pub triangular_inverse: Rate,
334 pub sparse_matmul: Rate,
337 pub householder_rank_one: Rate,
339 pub lu_swap: Rate,
343}
344
345const fn fmin(a: f64, b: f64) -> f64 {
346 if b < a { b } else { a }
347}
348
349const fn rmin(a: Rate, b: Rate) -> Rate {
350 Rate { real: fmin(a.real, b.real), cplx: fmin(a.cplx, b.cplx) }
351}
352
353const fn r(real: f64, cplx: f64) -> Rate {
354 Rate { real, cplx }
355}
356
357impl Rates {
358 const fn min_of(a: &Rates, b: &Rates) -> Rates {
360 Rates {
361 gemm_mac: rmin(a.gemm_mac, b.gemm_mac),
362 gemm_single: rmin(a.gemm_single, b.gemm_single),
363 block_apply_mac: rmin(a.block_apply_mac, b.block_apply_mac),
364 lu_full_pivot: rmin(a.lu_full_pivot, b.lu_full_pivot),
365 bunch_kaufman: rmin(a.bunch_kaufman, b.bunch_kaufman),
366 qr_col_pivot: rmin(a.qr_col_pivot, b.qr_col_pivot),
367 bidiag: rmin(a.bidiag, b.bidiag),
368 hessenberg: rmin(a.hessenberg, b.hessenberg),
369 tridiag: rmin(a.tridiag, b.tridiag),
370 matvec: rmin(a.matvec, b.matvec),
371 bidiag_svd_dc: rmin(a.bidiag_svd_dc, b.bidiag_svd_dc),
372 tridiag_evd_dc: rmin(a.tridiag_evd_dc, b.tridiag_evd_dc),
373 triangular_inverse: rmin(a.triangular_inverse, b.triangular_inverse),
374 sparse_matmul: rmin(a.sparse_matmul, b.sparse_matmul),
375 householder_rank_one: rmin(a.householder_rank_one, b.householder_rank_one),
376 lu_swap: rmin(a.lu_swap, b.lu_swap),
377 }
378 }
379
380 pub const X86_64_AVX512: Rates = Rates {
382 gemm_mac: r(0.0641, 0.2565),
383 gemm_single: r(0.0320, 0.1283),
384 block_apply_mac: r(0.0676, 0.2685),
385 lu_full_pivot: r(0.4011, 0.8679),
386 bunch_kaufman: r(0.8295, 1.3904),
387 qr_col_pivot: r(0.3645, 1.0320),
388 bidiag: r(1.0427, 2.1997),
389 hessenberg: r(0.8576, 2.2212),
390 tridiag: r(0.7695, 1.9252),
391 matvec: r(0.1654, 0.4215),
393 bidiag_svd_dc: r(152.9, 152.9),
394 tridiag_evd_dc: r(97.4, 97.4),
395 triangular_inverse: r(0.01936, 0.07228),
396 sparse_matmul: r(4.1315, 5.2342),
398 householder_rank_one: r(0.4783, 1.0303),
399 lu_swap: r(0.98, 1.30),
403 };
404
405 pub const AARCH64_APPLE: Rates = Rates {
407 gemm_mac: r(0.037, 0.165),
408 gemm_single: r(0.037 / 2.0, 0.165 / 2.0),
409 block_apply_mac: r(0.037, 0.165),
410 lu_full_pivot: r(0.216, 0.6),
411 bunch_kaufman: r(0.238, 0.38),
412 qr_col_pivot: r(0.202, 0.713),
413 bidiag: r(0.48, 1.36),
414 hessenberg: r(0.451, 1.57),
415 tridiag: r(0.339, 1.16),
416 matvec: r(0.1, 0.34),
417 bidiag_svd_dc: r(53.0, 53.0),
418 tridiag_evd_dc: r(31.0, 31.0),
419 triangular_inverse: r(0.0096, 0.0344),
420 sparse_matmul: r(1.5, 1.8),
421 householder_rank_one: r(0.212, 0.725),
422 lu_swap: r(NS_PER_SWAP, NS_PER_SWAP),
423 };
424
425 pub const GENERIC: Rates = Rates::min_of(&Rates::X86_64_AVX512, &Rates::AARCH64_APPLE);
427
428 pub fn for_class(class: HostClass) -> &'static Rates {
430 static AVX512: Rates = Rates::X86_64_AVX512;
431 static APPLE: Rates = Rates::AARCH64_APPLE;
432 static GENERIC: Rates = Rates::GENERIC;
433 match class {
434 HostClass::X86_64Avx512 => &AVX512,
435 HostClass::Aarch64Apple => &APPLE,
436 HostClass::X86_64Avx2 | HostClass::X86_64Scalar | HostClass::Aarch64 | HostClass::Generic => &GENERIC,
437 }
438 }
439}
440
441pub fn rates() -> &'static Rates {
443 Rates::for_class(HostClass::detect())
444}
445
446pub mod rate {
449 use super::{Rate, rates};
450 macro_rules! families {
451 ($($name:ident),* $(,)?) => {$(
452 #[doc = core::concat!("`Rates::", core::stringify!($name), "` of the host class.")]
453 pub fn $name() -> Rate { rates().$name }
454 )*};
455 }
456 families!(
457 gemm_mac,
458 block_apply_mac,
459 lu_full_pivot,
460 bunch_kaufman,
461 qr_col_pivot,
462 bidiag,
463 hessenberg,
464 tridiag,
465 matvec,
466 bidiag_svd_dc,
467 tridiag_evd_dc,
468 triangular_inverse,
469 sparse_matmul,
470 householder_rank_one,
471 lu_swap,
472 );
473}
474
475pub(crate) fn gemm_ns_per_mac<T: crate::ComplexField>() -> f64 {
479 if const { T::IS_NATIVE_F32 } {
480 ns_per_mac::f32()
481 } else if const { T::IS_NATIVE_C32 } {
482 ns_per_mac::c32()
483 } else if const { T::IS_REAL } {
484 ns_per_mac::f64()
485 } else {
486 ns_per_mac::c64()
487 }
488}
489
490
491#[cfg(feature = "par-trace")]
495pub mod trace {
496 use std::sync::Mutex;
497 use std::sync::atomic::{AtomicBool, Ordering};
498 #[derive(Clone, Debug)]
500 pub struct Record {
501 pub site: &'static str,
503 pub dims: [usize; 3],
505 pub units: usize,
507 pub sections: u32,
509 pub estimate_ns: f64,
511 pub dispatch_ns: f64,
513 pub width: usize,
515 pub parallel: bool,
517 }
518 pub static N_TASKS_MODE: core::sync::atomic::AtomicUsize = core::sync::atomic::AtomicUsize::new(0);
521 static ON: AtomicBool = AtomicBool::new(false);
522 static LOG: Mutex<Vec<Record>> = Mutex::new(Vec::new());
523 std::thread_local! {
524 static DIMS: core::cell::Cell<(&'static str, [usize; 3])> = const { core::cell::Cell::new(("", [0; 3])) };
525 }
526 pub fn enable(on: bool) {
528 if on {
529 LOG.lock().unwrap().clear();
530 }
531 ON.store(on, Ordering::SeqCst);
532 }
533 pub fn take() -> Vec<Record> {
535 core::mem::take(&mut *LOG.lock().unwrap())
536 }
537 pub fn shape(label: &'static str, a: usize, b: usize, c: usize) {
539 DIMS.with(|d| d.set((label, [a, b, c])));
540 }
541 pub(crate) fn record(
542 loc: &'static core::panic::Location<'static>,
543 units: usize,
544 sections: u32,
545 estimate_ns: f64,
546 dispatch_ns: f64,
547 width: usize,
548 parallel: bool,
549 ) {
550 if !ON.load(Ordering::Relaxed) {
551 return;
552 }
553 let (label, dims) = DIMS.with(|d| d.replace(("", [0; 3])));
554 let site: &'static str = if label.is_empty() {
555 Box::leak(format!("{}:{}", loc.file().rsplit('/').next().unwrap(), loc.line()).into_boxed_str())
556 } else {
557 label
558 };
559 LOG.lock().unwrap().push(Record { site, dims, units, sections, estimate_ns, dispatch_ns, width, parallel });
560 }
561}
562
563pub(crate) fn par_pays_time(par: &crate::Par, estimate: Duration, sections: u32, width: usize) -> bool {
568 let width = width.min(par.degree());
569 let Some(width) = NonZeroUsize::new(width).filter(|w| w.get() >= 2) else {
570 return false;
571 };
572 match par {
573 crate::Par::Seq => false,
574 #[cfg(feature = "threadteam")]
575 crate::Par::Team(team) => {
576 pays_at(estimate, sections, team.dispatch_cost(width))
579 },
580 #[cfg(not(feature = "threadteam"))]
582 _ => false,
583 }
584}
585
586pub(crate) fn par_can_pay_time(par: &crate::Par, estimate: Duration, sections: u32) -> bool {
591 if par.degree() < 2 {
592 return false;
593 }
594 match par {
595 crate::Par::Seq => false,
596 #[cfg(feature = "threadteam")]
597 crate::Par::Team(team) => par_can_pay_time_team(team, estimate, sections),
598 #[cfg(not(feature = "threadteam"))]
599 _ => false,
600 }
601}
602
603#[cfg(feature = "threadteam")]
605pub(crate) fn par_can_pay_time_team(team: &threadteam::TeamExec<'_>, estimate: Duration, sections: u32) -> bool {
606 team.size().get() >= 2
607 && pays_at(estimate, sections, team.dispatch_costs().min_cost(threadteam::DispatchKind::Team))
608}
609
610pub(crate) fn par_width_time(par: &crate::Par, estimate: Duration, sections: u32) -> usize {
618 par_width_time_upto(par, estimate, sections, usize::MAX)
619}
620
621pub(crate) fn par_width_time_upto(par: &crate::Par, estimate: Duration, sections: u32, max: usize) -> usize {
624 match par {
625 #[cfg(feature = "threadteam")]
626 crate::Par::Team(team) => par_width_team(team, estimate, sections, max),
627 _ => {
628 let _ = (estimate, sections, max);
629 1
630 },
631 }
632}
633
634#[cfg(feature = "threadteam")]
636pub(crate) fn par_width_team(team: &threadteam::TeamExec<'_>, estimate: Duration, sections: u32, max: usize) -> usize {
637 if max < 2 || !par_can_pay_time_team(team, estimate, sections) {
640 return 1;
641 }
642 (2..=team.size().get().min(max))
643 .rev()
644 .find(|&w| pays_at(estimate, sections, team.dispatch_cost(NonZeroUsize::new(w).unwrap())))
645 .unwrap_or(1)
646}
647
648#[cfg_attr(feature = "par-trace", track_caller)]
651pub(crate) fn par_width<T: crate::ComplexField>(
652 par: &crate::Par,
653 units: usize,
654 rate: Rate,
655 sections: u32,
656) -> usize {
657 par_width_upto::<T>(par, units, rate, sections, usize::MAX)
658}
659
660#[cfg_attr(feature = "par-trace", track_caller)]
662pub(crate) fn par_width_upto<T: crate::ComplexField>(
663 par: &crate::Par,
664 units: usize,
665 rate: Rate,
666 sections: u32,
667 max: usize,
668) -> usize {
669 if par.degree() < 2 || max < 2 {
670 return 1;
671 }
672 let w = par_width_time_upto(par, serial_time(units, rate.of::<T>()), sections, max);
673 #[cfg(feature = "par-trace")]
674 trace_decision(par.shape(), w, units, rate.of::<T>(), sections, w > 1, core::panic::Location::caller());
675 w
676}
677
678#[cfg(feature = "par-trace")]
679fn trace_decision(
680 shape: ParShape,
681 width: usize,
682 units: usize,
683 ns: f64,
684 sections: u32,
685 r: bool,
686 loc: &'static core::panic::Location<'static>,
687) {
688 let d = NonZeroUsize::new(width)
689 .map(|w| shape.dispatch_cost(w).as_secs_f64() * 1e9)
690 .unwrap_or(f64::INFINITY);
691 trace::record(loc, units, sections, units as f64 * ns, d, width, r);
692}
693
694#[cfg_attr(feature = "par-trace", track_caller)]
699pub(crate) fn join_pays<T: crate::ComplexField>(
700 par: &crate::Par,
701 units: usize,
702 rate: Rate,
703 sections: u32,
704) -> bool {
705 if par.degree() < 2 {
706 return false;
707 }
708 let r = par_pays_time(par, serial_time(units, rate.of::<T>()), sections, 2);
709 #[cfg(feature = "par-trace")]
710 trace_decision(par.shape(), 2, units, rate.of::<T>(), sections, r, core::panic::Location::caller());
711 r
712}
713
714pub(crate) fn halves_pays<T: crate::ComplexField>(
722 par: &crate::Par,
723 n: usize,
724 rate: Rate,
725 power: u32,
726) -> bool {
727 let h = n / 2;
728 join_pays::<T>(par, h.saturating_pow(power).saturating_mul(2), rate, 1)
729}
730
731pub(crate) fn block_apply_estimate<T: crate::ComplexField>(
738 m: usize,
739 n: usize,
740 k: usize,
741) -> Duration {
742 block_apply_estimate_with::<T>(rates(), m, n, k)
743}
744
745pub(crate) fn block_apply_estimate_with<T: crate::ComplexField>(
747 rt: &Rates,
748 m: usize,
749 n: usize,
750 k: usize,
751) -> Duration {
752 let macs = m.saturating_mul(n).saturating_mul(k).saturating_mul(2);
753 let ns = macs as f64 * rt.block_apply_mac.of::<T>() + m.saturating_mul(k) as f64 * rt.householder_rank_one.of::<T>();
755 Duration::from_secs_f64(ns * 1e-9)
756}
757
758#[derive(Copy, Clone, Debug)]
762pub(crate) struct PanelTeam {
763 pub from: usize,
764 pub to: usize,
765 pub ratio: f64,
766}
767
768pub(crate) fn lookahead_width(par: &crate::Par, panel: Duration, far: Duration, team: PanelTeam) -> Option<usize> {
782 let p = par.degree();
783 if p < 2 {
784 return None;
785 }
786 let (panel, far) = (panel.as_secs_f64(), far.as_secs_f64());
787 let panel_on = |a: usize| if (team.from..=team.to).contains(&a) { panel * team.ratio } else { panel };
788 let sequential = panel_on(p) + far / p as f64;
789 let split = par.shape().dispatch_cost(NonZeroUsize::new(p).unwrap()).as_secs_f64();
790 let mut best: Option<(usize, f64)> = None;
791 for a in [1usize, 4] {
792 if a >= p || (a > 1 && p < 8) {
793 continue;
794 }
795 let overlapped = f64::max(panel_on(a), far / (p - a) as f64) + split;
796 if best.is_none_or(|(_, c)| overlapped < c) {
797 best = Some((a, overlapped));
798 }
799 }
800 match best {
801 Some((a, c)) if c < sequential => Some(a),
802 _ => None,
803 }
804}
805
806pub(crate) fn serial_time(count: usize, ns: f64) -> Duration {
808 Duration::from_secs_f64(count as f64 * ns * 1e-9)
809}
810
811impl ParShape {
812 pub fn dispatch_cost(self, width: NonZeroUsize) -> Duration {
821 let width = width.min(NonZeroUsize::new(self.degree().max(1)).unwrap());
822 if width.get() < 2 {
823 return Duration::MAX;
824 }
825 match self {
826 Self::Seq => Duration::MAX,
827 #[cfg(feature = "threadteam")]
828 Self::Team(_, costs) => costs.cost(threadteam::DispatchKind::Team, width),
829 }
830 }
831
832 pub fn min_dispatch_cost(self) -> Duration {
836 if self.degree() < 2 {
837 return Duration::MAX;
838 }
839 match self {
840 Self::Seq => Duration::MAX,
841 #[cfg(feature = "threadteam")]
842 Self::Team(_, costs) => costs.min_cost(threadteam::DispatchKind::Team),
843 }
844 }
845
846 pub fn pays(self, estimate: Duration, sections: u32, width: NonZeroUsize) -> bool {
850 pays_at(estimate, sections, self.dispatch_cost(width))
851 }
852
853 pub fn pays_full(self, estimate: Duration, sections: u32) -> bool {
855 match NonZeroUsize::new(self.degree()) {
856 Some(width) => self.pays(estimate, sections, width),
857 None => false,
858 }
859 }
860}
861
862#[cfg(test)]
863mod tests {
864 #[cfg(feature = "threadteam")]
869 #[test]
870 fn lookahead_width_follows_the_step_model() {
871 use core::time::Duration;
872 let ms = Duration::from_millis;
873 let team = super::PanelTeam { from: 4, to: 32, ratio: 0.68 };
874 assert!(super::lookahead_width(&crate::Par::Seq, ms(5), ms(500), team).is_none());
875 crate::team::with_team(4, |par| {
876 assert_eq!(super::lookahead_width(&par, ms(5), ms(1), team), None);
878 assert_eq!(super::lookahead_width(&par, ms(5), ms(40), team), Some(1));
880 });
881 crate::team::with_team(16, |par| {
882 assert_eq!(super::lookahead_width(&par, ms(5), ms(400), team), Some(1));
884 });
885 crate::team::with_team(64, |par| {
886 assert_eq!(super::lookahead_width(&par, ms(20), ms(1000), team), Some(4));
889 });
890 }
891
892 use super::*;
893 use crate::c64;
894
895 fn nz(n: usize) -> NonZeroUsize {
896 NonZeroUsize::new(n).unwrap()
897 }
898
899 #[cfg(feature = "rayon")]
900 #[test]
901 fn seq_never_pays() {
902 assert_eq!(ParShape::Seq.dispatch_cost(nz(4)), Duration::MAX);
903 assert!(!ParShape::Seq.pays(Duration::from_secs(10), 1, nz(4)));
904 assert!(!ParShape::Seq.pays_full(Duration::from_secs(10), 1));
905 }
906
907 #[cfg(feature = "rayon")]
908 #[test]
909 fn min_dispatch_cost_is_the_cheapest_width() {
910 use threadteam::{DispatchCosts, DispatchKind, WidthCosts};
911 let us = Duration::from_micros;
912 for t in [[us(9), us(5), us(3)], [us(4), us(4), us(4)], [us(2), us(7), us(12)]] {
914 let costs = DispatchCosts::for_class(threadteam::HostClass::Generic).with(DispatchKind::Team, WidthCosts::new(t));
915 let s = ParShape::Team(nz(6), costs);
916 let want = (2..=6).map(|w| s.dispatch_cost(nz(w))).min().unwrap();
917 assert_eq!(s.min_dispatch_cost(), want, "{t:?}");
918 }
919 assert_eq!(ParShape::Seq.min_dispatch_cost(), Duration::MAX);
920 assert_eq!(ParShape::Team(nz(1), DispatchCosts::for_class(threadteam::HostClass::Generic)).min_dispatch_cost(), Duration::MAX);
921 }
922
923 #[cfg(feature = "rayon")]
924 #[test]
925 fn pays_is_the_multiple_of_the_cost() {
926 let d = Duration::from_micros(2);
927 assert!(pays_at(Duration::from_micros(6), 1, d));
928 assert!(!pays_at(Duration::from_nanos(5999), 1, d));
929 assert!(!pays_at(Duration::from_micros(11), 2, d));
930 assert!(pays_at(Duration::from_micros(12), 2, d));
931 assert!(!pays_at(Duration::from_secs(1), 1, Duration::MAX));
932 }
933
934 #[test]
935 fn team_shape_carries_the_costs_of_its_token() {
936 use threadteam::{DispatchCosts, DispatchKind, WidthCosts};
937 let costs = DispatchCosts::for_class(threadteam::HostClass::Generic)
938 .with(DispatchKind::Team, WidthCosts::flat(Duration::from_micros(5)));
939 let s = ParShape::Team(nz(4), costs);
940 assert_eq!(s.dispatch_cost(nz(4)), Duration::from_micros(5));
941 assert_eq!(s.dispatch_cost(nz(1)), Duration::MAX);
942 assert!(s.pays(Duration::from_micros(16), 1, nz(4)));
945 assert!(!s.pays(Duration::from_micros(14), 1, nz(4)));
946 }
947
948 #[test]
949 fn gemm_rates_are_per_dtype() {
950 assert_eq!(gemm_ns_per_mac::<f64>(), ns_per_mac::f64());
951 assert_eq!(gemm_ns_per_mac::<crate::c64>(), ns_per_mac::c64());
952 assert_eq!(gemm_ns_per_mac::<f32>(), ns_per_mac::f32());
953 assert_eq!(gemm_ns_per_mac::<crate::c32>(), ns_per_mac::c32());
954 assert!(ns_per_mac::c64() > 4.0 * ns_per_mac::f64());
956 }
957
958 fn floor_units<T: crate::ComplexField>(
960 shape: ParShape,
961 r: Rate,
962 sections: u32,
963 ) -> usize {
964 let pays = |u: usize| {
965 shape.pays_full(serial_time(u, r.of::<T>()), sections)
966 };
967 let mut hi = 1usize;
968 while !pays(hi) {
969 hi *= 2;
970 assert!(hi < 1 << 40, "never pays");
971 }
972 let mut lo = hi / 2;
973 while hi - lo > 1 {
974 let mid = (lo + hi) / 2;
975 if pays(mid) {
976 hi = mid;
977 } else {
978 lo = mid;
979 }
980 }
981 hi
982 }
983
984 fn gates(rt: &Rates) -> [(&'static str, Rate, u32); 10] {
986 [
987 ("lu_full_pivot", rt.lu_full_pivot, 1),
988 ("bunch_kaufman", rt.bunch_kaufman, 1),
989 ("qr_col_pivot", rt.qr_col_pivot, 1),
990 ("bidiag", rt.bidiag, 2),
991 ("hessenberg", rt.hessenberg, 2),
992 ("tridiag", rt.tridiag, 1),
993 ("matvec", rt.matvec, 1),
994 ("householder_rank_one", rt.householder_rank_one, 1),
995 ("gemm_mac", rt.gemm_mac, 1),
996 ("sparse_matmul", rt.sparse_matmul, 1),
997 ]
998 }
999
1000 const M5: &Rates = &Rates::AARCH64_APPLE;
1003
1004 fn entries(rt: &Rates) -> [(&'static str, Rate); 16] {
1006 [
1007 ("gemm_mac", rt.gemm_mac),
1008 ("gemm_single", rt.gemm_single),
1009 ("block_apply_mac", rt.block_apply_mac),
1010 ("lu_full_pivot", rt.lu_full_pivot),
1011 ("bunch_kaufman", rt.bunch_kaufman),
1012 ("qr_col_pivot", rt.qr_col_pivot),
1013 ("bidiag", rt.bidiag),
1014 ("hessenberg", rt.hessenberg),
1015 ("tridiag", rt.tridiag),
1016 ("matvec", rt.matvec),
1017 ("bidiag_svd_dc", rt.bidiag_svd_dc),
1018 ("tridiag_evd_dc", rt.tridiag_evd_dc),
1019 ("triangular_inverse", rt.triangular_inverse),
1020 ("sparse_matmul", rt.sparse_matmul),
1021 ("householder_rank_one", rt.householder_rank_one),
1022 ("lu_swap", rt.lu_swap),
1023 ]
1024 }
1025
1026 #[test]
1027 fn every_class_resolves_to_its_table() {
1028 assert_eq!(Rates::for_class(HostClass::X86_64Avx512), &Rates::X86_64_AVX512);
1029 assert_eq!(Rates::for_class(HostClass::Aarch64Apple), &Rates::AARCH64_APPLE);
1030 for c in [HostClass::X86_64Avx2, HostClass::X86_64Scalar, HostClass::Aarch64, HostClass::Generic] {
1031 assert_eq!(Rates::for_class(c), &Rates::GENERIC, "{}", c.name());
1032 }
1033 assert_eq!(Rates::X86_64_AVX512.gemm_mac, Rate { real: 0.0641, cplx: 0.2565 });
1035 assert_eq!(Rates::X86_64_AVX512.gemm_single, Rate { real: 0.0320, cplx: 0.1283 });
1036 assert_eq!(Rates::AARCH64_APPLE.gemm_mac, Rate { real: 0.037, cplx: 0.165 });
1037 assert_eq!(Rates::AARCH64_APPLE.bidiag_svd_dc, Rate { real: 53.0, cplx: 53.0 });
1038 let g = Rates::X86_64_AVX512;
1040 assert!((g.gemm_single.real / (g.gemm_mac.real / 2.0) - 1.0).abs() < 0.01);
1041 assert!((g.gemm_single.cplx / (g.gemm_mac.cplx / 2.0) - 1.0).abs() < 0.01);
1042 }
1043
1044 #[test]
1045 fn generic_is_the_per_entry_minimum_of_the_measured_rows() {
1046 let (a, b, g) = (entries(&Rates::X86_64_AVX512), entries(&Rates::AARCH64_APPLE), entries(&Rates::GENERIC));
1047 for i in 0..a.len() {
1048 let (name, ga) = (a[i].0, a[i].1);
1049 assert_eq!(g[i].1.real, ga.real.min(b[i].1.real), "{name} real");
1050 assert_eq!(g[i].1.cplx, ga.cplx.min(b[i].1.cplx), "{name} cplx");
1051 assert!(g[i].1.real <= ga.real && g[i].1.real <= b[i].1.real, "{name}");
1052 }
1053 }
1054
1055 #[test]
1056 fn rows_are_positive() {
1057 for c in HostClass::ALL {
1058 for (name, r) in entries(Rates::for_class(c)) {
1059 assert!(r.real > 0.0 && r.cplx > 0.0, "{} {name}", c.name());
1060 }
1061 }
1062 }
1063
1064 #[test]
1065 fn detection_is_stable_and_the_host_row_is_its_class() {
1066 let c = HostClass::detect();
1067 for _ in 0..4 {
1068 assert_eq!(HostClass::detect(), c);
1069 }
1070 assert_eq!(HostClass::detect_uncached(), c);
1071 assert_eq!(rates(), Rates::for_class(c));
1072 assert_eq!(rate::bidiag(), rates().bidiag);
1073 #[cfg(all(target_arch = "x86_64", feature = "std"))]
1074 if std::arch::is_x86_feature_detected!("avx512f") {
1075 assert_eq!(c, HostClass::X86_64Avx512);
1076 }
1077 #[cfg(all(target_arch = "aarch64", target_os = "macos"))]
1078 assert_eq!(c, HostClass::Aarch64Apple);
1079 }
1080
1081 #[test]
1083 fn host_class_agrees_with_threadteam() {
1084 let tt = |c: HostClass| match c {
1085 HostClass::X86_64Avx512 => threadteam::HostClass::X86_64Avx512,
1086 HostClass::X86_64Avx2 => threadteam::HostClass::X86_64Avx2,
1087 HostClass::X86_64Scalar => threadteam::HostClass::X86_64Scalar,
1088 HostClass::Aarch64Apple => threadteam::HostClass::Aarch64Apple,
1089 HostClass::Aarch64 => threadteam::HostClass::Aarch64,
1090 HostClass::Generic => threadteam::HostClass::Generic,
1091 };
1092 assert_eq!(tt(HostClass::detect()), threadteam::HostClass::detect());
1093 for c in HostClass::ALL {
1094 assert_eq!(c.name(), tt(c).name());
1095 }
1096 }
1097
1098 #[test]
1102 fn the_team_cost_is_threadteams_table_at_every_width() {
1103 use threadteam::{DispatchCosts, DispatchKind, HostClass, WidthCosts};
1104 let us = Duration::from_micros;
1105 let rows = [
1106 DispatchCosts::for_class(HostClass::Generic),
1107 DispatchCosts::for_class(HostClass::X86_64Avx512),
1108 DispatchCosts::for_class(HostClass::X86_64Avx2),
1109 DispatchCosts::for_class(HostClass::Aarch64Apple),
1110 DispatchCosts::for_class(HostClass::Generic).with(DispatchKind::Team, WidthCosts::flat(us(1))),
1111 ];
1112 for costs in rows {
1113 let s = ParShape::Team(nz(128), costs);
1114 for w in 2..=128 {
1115 assert_eq!(s.dispatch_cost(nz(w)), costs.cost(DispatchKind::Team, nz(w)), "{costs:?} width {w}");
1116 }
1117 }
1118 let flat = ParShape::Team(nz(128), rows[4]);
1121 for (w, ns) in [(8, 1333), (12, 2000), (16, 2666), (32, 5333), (64, 10666), (128, 21333)] {
1122 assert_eq!(flat.dispatch_cost(nz(w)), Duration::from_nanos(ns), "flat width {w}");
1123 }
1124 let avx512 = ParShape::Team(nz(128), rows[1]);
1125 assert_eq!(avx512.dispatch_cost(nz(8)), Duration::from_nanos(980));
1126 assert_eq!(avx512.dispatch_cost(nz(64)), Duration::from_nanos(980 * 8));
1127 let zen2 = ParShape::Team(nz(128), rows[2]);
1128 assert_eq!(zen2.dispatch_cost(nz(64)), Duration::from_nanos(9333));
1129 assert_eq!(zen2.dispatch_cost(nz(128)), Duration::from_nanos(18666));
1130 let at_8 = flat.dispatch_cost(nz(8)) * 3;
1132 assert!(flat.pays(at_8, 1, nz(8)));
1133 assert!(!flat.pays(at_8, 1, nz(64)));
1134 }
1135
1136 #[test]
1137 fn fixed_and_never_keep_their_meaning() {
1138 let r = rate::lu_full_pivot();
1139 assert_eq!(ParThreshold::Fixed(100).region_width::<f64>(&crate::Par::Seq, 100, r, 1), 1);
1141 crate::team::with_team(6, |par| {
1142 assert_eq!(ParThreshold::Never.region_width::<f64>(&par, usize::MAX, r, 1), 1);
1143 assert_eq!(ParThreshold::Fixed(100).region_width::<f64>(&par, 99, r, 1), 1);
1144 assert_eq!(ParThreshold::Fixed(100).region_width::<f64>(&par, 100, r, 1), 6);
1145 });
1146 assert_eq!(ParThreshold::from(0), ParThreshold::Derived);
1148 assert_eq!(ParThreshold::from(usize::MAX), ParThreshold::Never);
1149 assert_eq!(ParThreshold::from(7), ParThreshold::Fixed(7));
1150 assert_eq!(usize::from(ParThreshold::Fixed(0)), 1);
1151 }
1152
1153 #[test]
1154 fn derived_gates_never_pay_on_a_sequential_shape() {
1155 let seq = crate::Par::Seq;
1156 for (_, r, sections) in gates(rates()) {
1157 assert_eq!(ParThreshold::Derived.region_width::<f64>(&seq, usize::MAX / 2, r, sections), 1);
1158 assert_eq!(ParThreshold::Derived.region_width::<c64>(&seq, usize::MAX / 2, r, sections), 1);
1159 assert!(!ParShape::Seq.pays_full(serial_time(usize::MAX / 2, r.of::<f64>()), sections));
1160 }
1161 }
1162
1163 #[test]
1164 fn derived_floors_on_a_team_follow_the_rule() {
1165 use threadteam::DispatchCosts;
1166 let team = ParShape::Team(nz(6), DispatchCosts::for_class(threadteam::HostClass::Aarch64Apple));
1167 let d = team.dispatch_cost(nz(6)).as_secs_f64() * 1e9;
1168 for (name, r, sections) in gates(M5) {
1169 for (ty, ns) in [("f64", r.real), ("c64", r.cplx)] {
1170 let floor = if ty == "f64" {
1171 floor_units::<f64>(team, r, sections)
1172 } else {
1173 floor_units::<c64>(team, r, sections)
1174 };
1175 let want = PAYS * f64::from(sections) * d / ns;
1177 assert!(
1178 (floor as f64 - want).abs() <= 4.0 + want * 2e-3,
1179 "{name} {ty}: floor {floor}, rule {want}"
1180 );
1181 }
1182 }
1183 let pays = |u: usize, r: Rate, f: u32| {
1185 team.pays_full(serial_time(u, r.of::<f64>()), f)
1186 };
1187 assert!(!pays(100 * 100, M5.lu_full_pivot, 1));
1189 assert!(pays(200 * 200, M5.lu_full_pivot, 1));
1190 assert!(floor_units::<f64>(team, M5.bidiag, 2) >= 2 * floor_units::<f64>(team, M5.bidiag, 1) - 2);
1192 assert!(pays(256 * 256, M5.matvec, 1));
1194 assert!(!pays(64 * 64, M5.matvec, 1));
1195 let lu = floor_units::<f64>(team, M5.lu_full_pivot, 1);
1198 let bidiag = floor_units::<f64>(team, M5.bidiag, 1);
1199 assert!(lu > bidiag);
1200 }
1201
1202 #[test]
1203 fn par_width_reads_the_degree_of_the_token() {
1204 let seq = crate::Par::Seq;
1205 assert_eq!(par_width::<f64>(&seq, usize::MAX / 2, rate::matvec(), 1), 1);
1206 crate::team::with_team(6, |par| {
1207 assert_eq!(par_width::<f64>(&par, 10_000_000, rate::matvec(), 1), 6);
1208 assert_eq!(par_width::<f64>(&par, 1_000, rate::matvec(), 1), 1);
1209 assert_eq!(par_width_upto::<f64>(&par, 10_000_000, rate::matvec(), 1, 3), 3);
1211 });
1212 crate::team::with_team(1, |par| {
1213 assert_eq!(par_width::<f64>(&par, 10_000_000, rate::matvec(), 1), 1);
1214 });
1215 }
1216
1217 #[test]
1218 fn rates_are_per_dtype() {
1219 let r = rate::bidiag();
1220 assert_eq!(r.of::<f64>(), r.real);
1221 assert_eq!(r.of::<c64>(), r.cplx);
1222 assert_eq!(r.of::<f32>(), r.real / 2.0);
1223 assert_eq!(r.of::<crate::c32>(), r.cplx / 2.0);
1224 assert_eq!(rate::gemm_mac().of::<f64>(), ns_per_mac::f64());
1225 assert_eq!(rate::gemm_mac().of::<c64>(), ns_per_mac::c64());
1226 }
1227
1228 #[cfg(feature = "rayon")]
1232 mod forced_parallel {
1233 use super::*;
1234 use crate::dyn_stack::{MemBuffer, MemStack};
1235 use crate::internal_prelude::*;
1236 use crate::linalg::cholesky::lblt::factor as bk;
1237 use crate::linalg::evd::{hessenberg, tridiag};
1238 use crate::linalg::lu::full_pivoting::factor as fplu;
1239 use crate::linalg::qr::col_pivoting::factor as cpqr;
1240 use crate::linalg::svd::bidiag;
1241 use crate::{Auto, Mat, MatRef, Par};
1242
1243 const N: usize = 96;
1244
1245 fn fill<T: crate::ComplexField>(n: usize, seed: usize) -> Mat<T> {
1246 Mat::from_fn(n, n, |i, j| {
1247 let h = (i * 2654435761usize)
1248 .wrapping_add(j * 40503)
1249 .wrapping_add(seed * 977)
1250 % 1009;
1251 from_f64::<T>(h as f64 / 504.0 - 1.0)
1252 })
1253 }
1254 fn close<T: crate::ComplexField<Real = f64>>(
1255 what: &str,
1256 a: MatRef<'_, T>,
1257 b: MatRef<'_, T>,
1258 ) {
1259 let mut err = 0.0f64;
1260 for j in 0..a.ncols() {
1261 for i in 0..a.nrows() {
1262 err = err.max((a[(i, j)].clone() - b[(i, j)].clone()).abs());
1263 }
1264 }
1265 core::assert!(err <= 1e-9, "{what}: error {err}");
1266 }
1267
1268 fn check<T: crate::ComplexField<Real = f64>>(compare_complex_factors: bool) {
1269 let a0 = fill::<T>(N, 1);
1270 let sym = &a0 + a0.adjoint();
1271 let lu = |par: Par| {
1273 let mut a = a0.clone();
1274 let mut p = vec![0usize; 4 * N];
1275 let (r, rest) = p.split_at_mut(N);
1276 let (ri, rest) = rest.split_at_mut(N);
1277 let (c, ci) = rest.split_at_mut(N);
1278 let mut params: fplu::FullPivLuParams = Auto::<T>::auto();
1279 params.par_threshold = ParThreshold::Fixed(1);
1280 let mut mem = MemBuffer::new(fplu::lu_in_place_scratch::<usize, T>(
1281 N,
1282 N,
1283 par.shape(),
1284 params.into(),
1285 ));
1286 fplu::lu_in_place(
1287 a.as_mut(), r, ri, c, ci, par, MemStack::new(&mut mem), params.into(),
1288 );
1289 a
1290 };
1291 crate::team::with_team(4, |par| close("lu", lu(Par::Seq).as_ref(), lu(par).as_ref()));
1292 let bunch = |par: Par| {
1293 let mut a = sym.clone();
1294 let mut perm = vec![0usize; N];
1295 let mut perm_inv = vec![0usize; N];
1296 let mut sub = Mat::<T>::zeros(N, 1);
1297 let mut params: bk::LbltParams = Auto::<T>::auto();
1298 params.pivoting = bk::PivotingStrategy::Full;
1299 params.par_threshold = ParThreshold::Fixed(1);
1300 let mut mem = MemBuffer::new(bk::cholesky_in_place_scratch::<usize, T>(
1301 N,
1302 par.shape(),
1303 params.into(),
1304 ));
1305 bk::cholesky_in_place(
1306 a.as_mut(),
1307 sub.as_mut().col_mut(0).as_diagonal_mut(),
1308 &mut perm,
1309 &mut perm_inv,
1310 par,
1311 MemStack::new(&mut mem),
1312 params.into(),
1313 );
1314 a
1315 };
1316 crate::team::with_team(4, |par| close("bunch", bunch(Par::Seq).as_ref(), bunch(par).as_ref()));
1317 let qr = |par: Par| {
1318 let mut a = a0.clone();
1319 let mut h = Mat::<T>::zeros(16, N);
1320 let mut p = vec![0usize; N];
1321 let mut pi = vec![0usize; N];
1322 let mut params: cpqr::ColPivQrParams = Auto::<T>::auto();
1323 params.par_threshold = ParThreshold::Fixed(1);
1324 let mut mem = MemBuffer::new(cpqr::qr_in_place_scratch::<usize, T>(
1325 N,
1326 N,
1327 16,
1328 par.shape(),
1329 params.into(),
1330 ));
1331 cpqr::qr_in_place(
1332 a.as_mut(), h.as_mut(), &mut p, &mut pi, par, MemStack::new(&mut mem), params.into(),
1333 );
1334 a
1335 };
1336 crate::team::with_team(4, |par| close("qr", qr(Par::Seq).as_ref(), qr(par).as_ref()));
1337 let bid = |par: Par| {
1338 let mut a = a0.clone();
1339 let mut hl = Mat::<T>::zeros(8, N);
1340 let mut hr = Mat::<T>::zeros(8, N - 1);
1341 let mut params: bidiag::BidiagParams = Auto::<T>::auto();
1342 params.par_threshold = ParThreshold::Fixed(1);
1343 let mut mem = MemBuffer::new(bidiag::bidiag_in_place_scratch::<T>(
1344 N,
1345 N,
1346 par.shape(),
1347 params.into(),
1348 ));
1349 bidiag::bidiag_in_place(
1350 a.as_mut(), hl.as_mut(), hr.as_mut(), par, MemStack::new(&mut mem), params.into(),
1351 );
1352 for j in 0..N {
1357 for i in 0..N {
1358 if i > j || j > i + 1 {
1359 a[(i, j)] = T::zero();
1360 }
1361 }
1362 }
1363 a
1364 };
1365 if compare_complex_factors {
1372 crate::team::with_team(4, |par| {
1373 close("bid", bid(Par::Seq).as_ref(), bid(par).as_ref())
1374 });
1375 } else {
1376 let seq = bid(Par::Seq);
1380 let team = crate::team::with_team(4, |par| bid(par));
1381 core::assert!(seq.as_ref().is_all_finite(), "bid sequential");
1382 core::assert!(team.as_ref().is_all_finite(), "bid team");
1383 }
1384 let hess = |par: Par| {
1385 let mut a = a0.clone();
1386 let mut h = Mat::<T>::zeros(8, N - 1);
1387 let mut params: hessenberg::HessenbergParams = Auto::<T>::auto();
1388 params.par_threshold = ParThreshold::Fixed(1);
1389 params.blocking_threshold = usize::MAX;
1390 let mut mem = MemBuffer::new(hessenberg::hessenberg_in_place_scratch::<T>(
1391 N,
1392 8,
1393 par.shape(),
1394 params.into(),
1395 ));
1396 hessenberg::hessenberg_in_place(
1397 a.as_mut(), h.as_mut(), par, MemStack::new(&mut mem), params.into(),
1398 );
1399 a
1400 };
1401 if compare_complex_factors {
1406 crate::team::with_team(4, |par| {
1407 close("hess", hess(Par::Seq).as_ref(), hess(par).as_ref())
1408 });
1409 } else {
1410 let seq = hess(Par::Seq);
1411 let team = crate::team::with_team(4, |par| hess(par));
1412 core::assert!(seq.as_ref().is_all_finite(), "hess sequential");
1413 core::assert!(team.as_ref().is_all_finite(), "hess team");
1414 }
1415 let tri = |par: Par| {
1416 let mut a = sym.clone();
1417 let mut h = Mat::<T>::zeros(8, N - 1);
1418 let mut params: tridiag::TridiagParams = Auto::<T>::auto();
1419 params.par_threshold = ParThreshold::Fixed(1);
1420 let mut mem = MemBuffer::new(tridiag::tridiag_in_place_scratch::<T>(
1421 N,
1422 par.shape(),
1423 params.into(),
1424 ));
1425 tridiag::tridiag_in_place(
1426 a.as_mut(), h.as_mut(), par, MemStack::new(&mut mem), params.into(),
1427 );
1428 a
1429 };
1430 crate::team::with_team(4, |par| close("tri", tri(Par::Seq).as_ref(), tri(par).as_ref()));
1431 }
1432
1433 #[test]
1434 fn forced_parallel_steps_equal_serial_f64() {
1435 check::<f64>(true);
1436 }
1437
1438 #[test]
1439 fn forced_parallel_steps_equal_serial_c64() {
1440 check::<c64>(false);
1441 }
1442 }
1443
1444 #[test]
1446 fn halves_join_gates_follow_the_rule() {
1447 let seq = crate::Par::Seq;
1448 for r in [rate::bidiag_svd_dc(), rate::tridiag_evd_dc(), rate::triangular_inverse()] {
1449 assert!(!halves_pays::<f64>(&seq, 1 << 20, r, 2));
1450 }
1451 crate::team::with_team(6, |par| {
1452 for r in [rate::bidiag_svd_dc(), rate::tridiag_evd_dc(), rate::triangular_inverse()] {
1453 let power = if r == rate::triangular_inverse() { 3 } else { 2 };
1455 let first = (2..4096).find(|&n| halves_pays::<f64>(&par, n, r, power)).expect("never pays");
1456 assert!(first > 1, "{first}");
1459 assert!(!halves_pays::<f64>(&par, first - 1, r, power));
1460 assert!(halves_pays::<f64>(&par, first, r, power));
1461 assert!(halves_pays::<f64>(&par, 2 * first, r, power));
1462 let h = first / 2;
1464 let work_ns = 2.0 * (h as f64).powi(power as i32) * r.real;
1465 let want = PAYS * par.shape().dispatch_cost(nz(2)).as_secs_f64() * 1e9;
1466 assert!(work_ns >= want * 0.999, "{work_ns} {want}");
1467 }
1468 });
1469 }
1470
1471 #[test]
1478 fn block_apply_split_starts_where_the_estimate_pays() {
1479 use threadteam::DispatchCosts;
1480 let team = ParShape::Team(nz(6), DispatchCosts::for_class(threadteam::HostClass::Aarch64Apple));
1481 let pays = |m: usize| {
1482 let (n, k) = (8, m - 8);
1483 pays_at(block_apply_estimate_with::<f64>(M5, m, n, k), 1, team.dispatch_cost(nz(6)))
1484 };
1485 assert!(!pays(72));
1486 assert!(pays(80));
1487 let gemm_only = |m: usize| {
1489 pays_at(serial_time(2 * m * 8 * (m - 8), M5.gemm_mac.real), 1, team.dispatch_cost(nz(6)))
1490 };
1491 assert!(!gemm_only(88) && gemm_only(96));
1492 let c = block_apply_estimate_with::<c64>(M5, 100, 8, 92).as_secs_f64() * 1e9;
1494 let want = 2.0 * 100.0 * 8.0 * 92.0 * M5.block_apply_mac.cplx + 100.0 * 92.0 * M5.householder_rank_one.cplx;
1495 assert!((c - want).abs() < 2.0, "{c} {want}");
1496 }
1497
1498 #[test]
1499 fn join_gates_on_a_team_follow_the_width_two_cost() {
1500 use threadteam::DispatchCosts;
1501 let team = ParShape::Team(nz(6), DispatchCosts::for_class(threadteam::HostClass::Aarch64Apple));
1502 for (name, r, sections) in gates(M5) {
1503 assert!(floor_units::<f64>(team, r, sections) > 0, "{name}");
1504 }
1505 let d2 = team.dispatch_cost(nz(2));
1507 assert_eq!(d2, Duration::from_micros(3));
1508 let par = |ns: f64| pays_at(Duration::from_secs_f64(ns * 1e-9), 1, d2);
1509 assert!(par(9000.0) && !par(8900.0));
1510 for (r, power, yes, no) in [
1514 (M5.bidiag_svd_dc, 2u32, 128usize, 8usize),
1515 (M5.tridiag_evd_dc, 2, 64, 8),
1516 (M5.triangular_inverse, 3, 256, 64),
1517 ] {
1518 let pays = |n: usize| {
1519 let h = n / 2;
1520 pays_at(
1521 serial_time(2 * h.pow(power), r.real),
1522 1,
1523 d2,
1524 )
1525 };
1526 assert!(pays(yes));
1527 assert!(!pays(no));
1528 }
1529 assert!(team.pays_full(serial_time(1_000_000, M5.sparse_matmul.of::<f64>()), 1));
1531 assert!(!team.pays_full(serial_time(1_000, M5.sparse_matmul.of::<f64>()), 1));
1532 }
1533}
1534
1535#[cfg(feature = "par-trace")]
1536#[doc(hidden)]
1537pub fn trace_pub(
1538 units: usize,
1539 sections: u32,
1540 estimate: Duration,
1541 dispatch: Duration,
1542 width: usize,
1543 r: bool,
1544 loc: &'static core::panic::Location<'static>,
1545) {
1546 trace::record(loc, units, sections, estimate.as_secs_f64() * 1e9, dispatch.as_secs_f64() * 1e9, width, r);
1547}
1548
1549#[cfg(all(test, feature = "sparse-linalg"))]
1550mod join_results {
1551 use crate::dyn_stack::{MemBuffer, MemStack};
1554 use crate::linalg::evd::{self, ComputeEigenvectors};
1555 use crate::linalg::svd;
1556 use crate::linalg::triangular_inverse::{invert_lower_triangular, invert_unit_lower_triangular};
1557 use crate::{Accum, Mat, Par, Spec};
1558
1559 struct Rng(u64);
1560 impl Rng {
1561 fn next(&mut self) -> f64 {
1562 self.0 ^= self.0 << 13;
1563 self.0 ^= self.0 >> 7;
1564 self.0 ^= self.0 << 17;
1565 (self.0 >> 11) as f64 / (1u64 << 53) as f64 - 0.5
1566 }
1567 }
1568
1569 fn close(a: &[f64], b: &[f64], tol: f64) {
1570 let scale = a.iter().fold(1.0f64, |m, x| m.max(x.abs()));
1571 for (x, y) in a.iter().zip(b) {
1572 assert!((x - y).abs() <= tol * scale, "{x} {y}");
1573 }
1574 }
1575
1576 #[test]
1577 fn triangular_inverse_agrees_with_serial() {
1578 let n = 512;
1579 let mut r = Rng(3);
1580 let mut src = Mat::<f64>::from_fn(n, n, |_, _| r.next() * 0.1);
1581 for i in 0..n {
1582 src[(i, i)] = 2.0 + src[(i, i)].abs();
1583 }
1584 for unit in [false, true] {
1585 let mut a = Mat::<f64>::zeros(n, n);
1586 let mut b = Mat::<f64>::zeros(n, n);
1587 if unit {
1588 invert_unit_lower_triangular(a.as_mut(), src.as_ref(), Par::Seq);
1589 } else {
1590 invert_lower_triangular(a.as_mut(), src.as_ref(), Par::Seq);
1591 }
1592 crate::team::with_team(4, |par| {
1593 if unit {
1594 invert_unit_lower_triangular(b.as_mut(), src.as_ref(), par);
1595 } else {
1596 invert_lower_triangular(b.as_mut(), src.as_ref(), par);
1597 }
1598 });
1599 for j in 0..n {
1600 close(
1601 &(0..n).map(|i| a[(i, j)]).collect::<Vec<_>>(),
1602 &(0..n).map(|i| b[(i, j)]).collect::<Vec<_>>(),
1603 1e-9,
1604 );
1605 }
1606 }
1607 }
1608
1609 #[test]
1610 fn bidiagonal_svd_agrees_with_serial() {
1611 let n = 512;
1612 let mut r = Rng(5);
1613 let d0: Vec<f64> = (0..n).map(|_| r.next() + 1.0).collect();
1614 let e: Vec<f64> = (0..n - 1).map(|_| r.next()).collect();
1615 let mut out = vec![];
1616 let run = |par: Par, out: &mut Vec<Vec<f64>>| {
1617 let mut d = d0.clone();
1618 let mut u = Mat::<f64>::zeros(n, n);
1619 let mut v = Mat::<f64>::zeros(n, n);
1620 let mut mem = MemBuffer::new(svd::bidiagonal_svd_scratch::<f64>(
1621 n,
1622 par.shape(),
1623 Spec::default(),
1624 ));
1625 svd::bidiagonal_svd(
1626 &mut d,
1627 &e,
1628 true,
1629 u.as_mut(),
1630 v.as_mut(),
1631 par,
1632 MemStack::new(&mut mem),
1633 Spec::default(),
1634 )
1635 .unwrap();
1636 out.push(d);
1637 };
1638 run(Par::Seq, &mut out);
1639 crate::team::with_team(4, |par| run(par, &mut out));
1640 close(&out[0], &out[1], 1e-10);
1641 }
1642
1643 #[test]
1644 fn tridiagonal_evd_agrees_with_serial() {
1645 let n = 512;
1646 let mut r = Rng(6);
1647 let d0: Vec<f64> = (0..n).map(|_| r.next()).collect();
1648 let mut e0: Vec<f64> = (0..n - 1).map(|_| r.next()).collect();
1649 e0.push(0.0);
1650 let mut out = vec![];
1651 let run = |par: Par, out: &mut Vec<Vec<f64>>| {
1652 let mut u = Mat::<f64>::zeros(n, n);
1653 let mut s = Mat::<f64>::zeros(n, 1);
1654 let mut mem = MemBuffer::new(evd::tridiagonal_self_adjoint_evd_scratch::<f64>(
1655 n,
1656 ComputeEigenvectors::Yes,
1657 par.shape(),
1658 Spec::default(),
1659 ));
1660 let d = crate::col::ColRef::from_slice(&d0);
1661 let e = crate::col::ColRef::from_slice(&e0);
1662 evd::tridiagonal_self_adjoint_evd(
1663 d.as_diagonal(),
1664 e.as_diagonal(),
1665 s.as_mut().col_mut(0).as_diagonal_mut(),
1666 Some(u.as_mut()),
1667 par,
1668 MemStack::new(&mut mem),
1669 Spec::default(),
1670 )
1671 .unwrap();
1672 out.push((0..n).map(|i| s[(i, 0)]).collect::<Vec<_>>());
1673 };
1674 run(Par::Seq, &mut out);
1675 crate::team::with_team(4, |par| run(par, &mut out));
1676 close(&out[0], &out[1], 1e-10);
1677 }
1678
1679 #[test]
1680 fn sparse_matmul_agrees_with_serial() {
1681 use crate::sparse::SparseColMat;
1682 use crate::sparse::Triplet;
1683 use crate::sparse::linalg::matmul::sparse_sparse_matmul;
1684 let n = 6000;
1685 let mk = |seed: u64| {
1686 let mut r = Rng(seed);
1687 let mut t = vec![];
1688 for j in 0..n {
1689 for _ in 0..24 {
1690 let i = ((r.next() + 0.5) * n as f64) as usize % n;
1691 t.push(Triplet::new(i, j, r.next()));
1692 }
1693 }
1694 SparseColMat::<usize, f64>::try_new_from_triplets(n, n, &t).unwrap()
1695 };
1696 let (a, b) = (mk(1), mk(2));
1697 let s = sparse_sparse_matmul(a.as_ref(), b.as_ref(), 1.0, Par::Seq).unwrap();
1698 let p = crate::team::with_team(4, |par| {
1699 sparse_sparse_matmul(a.as_ref(), b.as_ref(), 1.0, par).unwrap()
1700 });
1701 close(s.val(), p.val(), 1e-12);
1702 let _ = Accum::Add;
1703 }
1704}