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 pub static LOWER_SPLIT: core::sync::atomic::AtomicUsize = core::sync::atomic::AtomicUsize::new(0);
524 static ON: AtomicBool = AtomicBool::new(false);
525 static LOG: Mutex<Vec<Record>> = Mutex::new(Vec::new());
526 std::thread_local! {
527 static DIMS: core::cell::Cell<(&'static str, [usize; 3])> = const { core::cell::Cell::new(("", [0; 3])) };
528 }
529 pub fn enable(on: bool) {
531 if on {
532 LOG.lock().unwrap().clear();
533 }
534 ON.store(on, Ordering::SeqCst);
535 }
536 pub fn take() -> Vec<Record> {
538 core::mem::take(&mut *LOG.lock().unwrap())
539 }
540 pub fn shape(label: &'static str, a: usize, b: usize, c: usize) {
542 DIMS.with(|d| d.set((label, [a, b, c])));
543 }
544 pub(crate) fn record(
545 loc: &'static core::panic::Location<'static>,
546 units: usize,
547 sections: u32,
548 estimate_ns: f64,
549 dispatch_ns: f64,
550 width: usize,
551 parallel: bool,
552 ) {
553 if !ON.load(Ordering::Relaxed) {
554 return;
555 }
556 let (label, dims) = DIMS.with(|d| d.replace(("", [0; 3])));
557 let site: &'static str = if label.is_empty() {
558 Box::leak(format!("{}:{}", loc.file().rsplit('/').next().unwrap(), loc.line()).into_boxed_str())
559 } else {
560 label
561 };
562 LOG.lock().unwrap().push(Record { site, dims, units, sections, estimate_ns, dispatch_ns, width, parallel });
563 }
564}
565
566pub(crate) fn par_pays_time(par: &crate::Par, estimate: Duration, sections: u32, width: usize) -> bool {
571 let width = width.min(par.degree());
572 let Some(width) = NonZeroUsize::new(width).filter(|w| w.get() >= 2) else {
573 return false;
574 };
575 match par {
576 crate::Par::Seq => false,
577 #[cfg(feature = "threadteam")]
578 crate::Par::Team(team) => {
579 pays_at(estimate, sections, team.dispatch_cost(width))
582 },
583 #[cfg(not(feature = "threadteam"))]
585 _ => false,
586 }
587}
588
589pub(crate) fn par_can_pay_time(par: &crate::Par, estimate: Duration, sections: u32) -> bool {
594 if par.degree() < 2 {
595 return false;
596 }
597 match par {
598 crate::Par::Seq => false,
599 #[cfg(feature = "threadteam")]
600 crate::Par::Team(team) => par_can_pay_time_team(team, estimate, sections),
601 #[cfg(not(feature = "threadteam"))]
602 _ => false,
603 }
604}
605
606#[cfg(feature = "threadteam")]
608pub(crate) fn par_can_pay_time_team(team: &threadteam::TeamExec<'_>, estimate: Duration, sections: u32) -> bool {
609 team.size().get() >= 2
610 && pays_at(estimate, sections, team.dispatch_costs().min_cost(threadteam::DispatchKind::Team))
611}
612
613pub(crate) fn par_width_time(par: &crate::Par, estimate: Duration, sections: u32) -> usize {
621 par_width_time_upto(par, estimate, sections, usize::MAX)
622}
623
624pub(crate) fn par_width_time_upto(par: &crate::Par, estimate: Duration, sections: u32, max: usize) -> usize {
627 match par {
628 #[cfg(feature = "threadteam")]
629 crate::Par::Team(team) => par_width_team(team, estimate, sections, max),
630 _ => {
631 let _ = (estimate, sections, max);
632 1
633 },
634 }
635}
636
637#[cfg(feature = "threadteam")]
639pub(crate) fn par_width_team(team: &threadteam::TeamExec<'_>, estimate: Duration, sections: u32, max: usize) -> usize {
640 if max < 2 || !par_can_pay_time_team(team, estimate, sections) {
643 return 1;
644 }
645 (2..=team.size().get().min(max))
646 .rev()
647 .find(|&w| pays_at(estimate, sections, team.dispatch_cost(NonZeroUsize::new(w).unwrap())))
648 .unwrap_or(1)
649}
650
651#[cfg_attr(feature = "par-trace", track_caller)]
654pub(crate) fn par_width<T: crate::ComplexField>(
655 par: &crate::Par,
656 units: usize,
657 rate: Rate,
658 sections: u32,
659) -> usize {
660 par_width_upto::<T>(par, units, rate, sections, usize::MAX)
661}
662
663#[cfg_attr(feature = "par-trace", track_caller)]
665pub(crate) fn par_width_upto<T: crate::ComplexField>(
666 par: &crate::Par,
667 units: usize,
668 rate: Rate,
669 sections: u32,
670 max: usize,
671) -> usize {
672 if par.degree() < 2 || max < 2 {
673 return 1;
674 }
675 let w = par_width_time_upto(par, serial_time(units, rate.of::<T>()), sections, max);
676 #[cfg(feature = "par-trace")]
677 trace_decision(par.shape(), w, units, rate.of::<T>(), sections, w > 1, core::panic::Location::caller());
678 w
679}
680
681#[cfg(feature = "par-trace")]
682fn trace_decision(
683 shape: ParShape,
684 width: usize,
685 units: usize,
686 ns: f64,
687 sections: u32,
688 r: bool,
689 loc: &'static core::panic::Location<'static>,
690) {
691 let d = NonZeroUsize::new(width)
692 .map(|w| shape.dispatch_cost(w).as_secs_f64() * 1e9)
693 .unwrap_or(f64::INFINITY);
694 trace::record(loc, units, sections, units as f64 * ns, d, width, r);
695}
696
697#[cfg_attr(feature = "par-trace", track_caller)]
702pub(crate) fn join_pays<T: crate::ComplexField>(
703 par: &crate::Par,
704 units: usize,
705 rate: Rate,
706 sections: u32,
707) -> bool {
708 if par.degree() < 2 {
709 return false;
710 }
711 let r = par_pays_time(par, serial_time(units, rate.of::<T>()), sections, 2);
712 #[cfg(feature = "par-trace")]
713 trace_decision(par.shape(), 2, units, rate.of::<T>(), sections, r, core::panic::Location::caller());
714 r
715}
716
717pub(crate) fn halves_pays<T: crate::ComplexField>(
725 par: &crate::Par,
726 n: usize,
727 rate: Rate,
728 power: u32,
729) -> bool {
730 let h = n / 2;
731 join_pays::<T>(par, h.saturating_pow(power).saturating_mul(2), rate, 1)
732}
733
734pub(crate) fn block_apply_estimate<T: crate::ComplexField>(
741 m: usize,
742 n: usize,
743 k: usize,
744) -> Duration {
745 block_apply_estimate_with::<T>(rates(), m, n, k)
746}
747
748pub(crate) fn block_apply_estimate_with<T: crate::ComplexField>(
750 rt: &Rates,
751 m: usize,
752 n: usize,
753 k: usize,
754) -> Duration {
755 let macs = m.saturating_mul(n).saturating_mul(k).saturating_mul(2);
756 let ns = macs as f64 * rt.block_apply_mac.of::<T>() + m.saturating_mul(k) as f64 * rt.householder_rank_one.of::<T>();
758 Duration::from_secs_f64(ns * 1e-9)
759}
760
761#[derive(Copy, Clone, Debug)]
765pub(crate) struct PanelTeam {
766 pub from: usize,
767 pub to: usize,
768 pub ratio: f64,
769}
770
771pub(crate) fn lookahead_width(par: &crate::Par, panel: Duration, far: Duration, team: PanelTeam) -> Option<usize> {
785 let p = par.degree();
786 if p < 2 {
787 return None;
788 }
789 let (panel, far) = (panel.as_secs_f64(), far.as_secs_f64());
790 let panel_on = |a: usize| if (team.from..=team.to).contains(&a) { panel * team.ratio } else { panel };
791 let sequential = panel_on(p) + far / p as f64;
792 let split = par.shape().dispatch_cost(NonZeroUsize::new(p).unwrap()).as_secs_f64();
793 let mut best: Option<(usize, f64)> = None;
794 for a in [1usize, 4] {
795 if a >= p || (a > 1 && p < 8) {
796 continue;
797 }
798 let overlapped = f64::max(panel_on(a), far / (p - a) as f64) + split;
799 if best.is_none_or(|(_, c)| overlapped < c) {
800 best = Some((a, overlapped));
801 }
802 }
803 match best {
804 Some((a, c)) if c < sequential => Some(a),
805 _ => None,
806 }
807}
808
809pub(crate) fn serial_time(count: usize, ns: f64) -> Duration {
811 Duration::from_secs_f64(count as f64 * ns * 1e-9)
812}
813
814impl ParShape {
815 pub fn dispatch_cost(self, width: NonZeroUsize) -> Duration {
824 let width = width.min(NonZeroUsize::new(self.degree().max(1)).unwrap());
825 if width.get() < 2 {
826 return Duration::MAX;
827 }
828 match self {
829 Self::Seq => Duration::MAX,
830 #[cfg(feature = "threadteam")]
831 Self::Team(_, costs) => costs.cost(threadteam::DispatchKind::Team, width),
832 }
833 }
834
835 pub fn min_dispatch_cost(self) -> Duration {
839 if self.degree() < 2 {
840 return Duration::MAX;
841 }
842 match self {
843 Self::Seq => Duration::MAX,
844 #[cfg(feature = "threadteam")]
845 Self::Team(_, costs) => costs.min_cost(threadteam::DispatchKind::Team),
846 }
847 }
848
849 pub fn pays(self, estimate: Duration, sections: u32, width: NonZeroUsize) -> bool {
853 pays_at(estimate, sections, self.dispatch_cost(width))
854 }
855
856 pub fn pays_full(self, estimate: Duration, sections: u32) -> bool {
858 match NonZeroUsize::new(self.degree()) {
859 Some(width) => self.pays(estimate, sections, width),
860 None => false,
861 }
862 }
863}
864
865#[cfg(test)]
866mod tests {
867 #[cfg(feature = "threadteam")]
872 #[test]
873 fn lookahead_width_follows_the_step_model() {
874 use core::time::Duration;
875 let ms = Duration::from_millis;
876 let team = super::PanelTeam { from: 4, to: 32, ratio: 0.68 };
877 assert!(super::lookahead_width(&crate::Par::Seq, ms(5), ms(500), team).is_none());
878 crate::team::with_team(4, |par| {
879 assert_eq!(super::lookahead_width(&par, ms(5), ms(1), team), None);
881 assert_eq!(super::lookahead_width(&par, ms(5), ms(40), team), Some(1));
883 });
884 crate::team::with_team(16, |par| {
885 assert_eq!(super::lookahead_width(&par, ms(5), ms(400), team), Some(1));
887 });
888 crate::team::with_team(64, |par| {
889 assert_eq!(super::lookahead_width(&par, ms(20), ms(1000), team), Some(4));
892 });
893 }
894
895 use super::*;
896 use crate::c64;
897
898 fn nz(n: usize) -> NonZeroUsize {
899 NonZeroUsize::new(n).unwrap()
900 }
901
902 #[cfg(feature = "rayon")]
903 #[test]
904 fn seq_never_pays() {
905 assert_eq!(ParShape::Seq.dispatch_cost(nz(4)), Duration::MAX);
906 assert!(!ParShape::Seq.pays(Duration::from_secs(10), 1, nz(4)));
907 assert!(!ParShape::Seq.pays_full(Duration::from_secs(10), 1));
908 }
909
910 #[cfg(feature = "rayon")]
911 #[test]
912 fn min_dispatch_cost_is_the_cheapest_width() {
913 use threadteam::{DispatchCosts, DispatchKind, WidthCosts};
914 let us = Duration::from_micros;
915 for t in [[us(9), us(5), us(3)], [us(4), us(4), us(4)], [us(2), us(7), us(12)]] {
917 let costs = DispatchCosts::for_class(threadteam::HostClass::Generic).with(DispatchKind::Team, WidthCosts::new(t));
918 let s = ParShape::Team(nz(6), costs);
919 let want = (2..=6).map(|w| s.dispatch_cost(nz(w))).min().unwrap();
920 assert_eq!(s.min_dispatch_cost(), want, "{t:?}");
921 }
922 assert_eq!(ParShape::Seq.min_dispatch_cost(), Duration::MAX);
923 assert_eq!(ParShape::Team(nz(1), DispatchCosts::for_class(threadteam::HostClass::Generic)).min_dispatch_cost(), Duration::MAX);
924 }
925
926 #[cfg(feature = "rayon")]
927 #[test]
928 fn pays_is_the_multiple_of_the_cost() {
929 let d = Duration::from_micros(2);
930 assert!(pays_at(Duration::from_micros(6), 1, d));
931 assert!(!pays_at(Duration::from_nanos(5999), 1, d));
932 assert!(!pays_at(Duration::from_micros(11), 2, d));
933 assert!(pays_at(Duration::from_micros(12), 2, d));
934 assert!(!pays_at(Duration::from_secs(1), 1, Duration::MAX));
935 }
936
937 #[test]
938 fn team_shape_carries_the_costs_of_its_token() {
939 use threadteam::{DispatchCosts, DispatchKind, WidthCosts};
940 let costs = DispatchCosts::for_class(threadteam::HostClass::Generic)
941 .with(DispatchKind::Team, WidthCosts::flat(Duration::from_micros(5)));
942 let s = ParShape::Team(nz(4), costs);
943 assert_eq!(s.dispatch_cost(nz(4)), Duration::from_micros(5));
944 assert_eq!(s.dispatch_cost(nz(1)), Duration::MAX);
945 assert!(s.pays(Duration::from_micros(16), 1, nz(4)));
948 assert!(!s.pays(Duration::from_micros(14), 1, nz(4)));
949 }
950
951 #[test]
952 fn gemm_rates_are_per_dtype() {
953 assert_eq!(gemm_ns_per_mac::<f64>(), ns_per_mac::f64());
954 assert_eq!(gemm_ns_per_mac::<crate::c64>(), ns_per_mac::c64());
955 assert_eq!(gemm_ns_per_mac::<f32>(), ns_per_mac::f32());
956 assert_eq!(gemm_ns_per_mac::<crate::c32>(), ns_per_mac::c32());
957 assert!(ns_per_mac::c64() > 4.0 * ns_per_mac::f64());
959 }
960
961 fn floor_units<T: crate::ComplexField>(
963 shape: ParShape,
964 r: Rate,
965 sections: u32,
966 ) -> usize {
967 let pays = |u: usize| {
968 shape.pays_full(serial_time(u, r.of::<T>()), sections)
969 };
970 let mut hi = 1usize;
971 while !pays(hi) {
972 hi *= 2;
973 assert!(hi < 1 << 40, "never pays");
974 }
975 let mut lo = hi / 2;
976 while hi - lo > 1 {
977 let mid = (lo + hi) / 2;
978 if pays(mid) {
979 hi = mid;
980 } else {
981 lo = mid;
982 }
983 }
984 hi
985 }
986
987 fn gates(rt: &Rates) -> [(&'static str, Rate, u32); 10] {
989 [
990 ("lu_full_pivot", rt.lu_full_pivot, 1),
991 ("bunch_kaufman", rt.bunch_kaufman, 1),
992 ("qr_col_pivot", rt.qr_col_pivot, 1),
993 ("bidiag", rt.bidiag, 2),
994 ("hessenberg", rt.hessenberg, 2),
995 ("tridiag", rt.tridiag, 1),
996 ("matvec", rt.matvec, 1),
997 ("householder_rank_one", rt.householder_rank_one, 1),
998 ("gemm_mac", rt.gemm_mac, 1),
999 ("sparse_matmul", rt.sparse_matmul, 1),
1000 ]
1001 }
1002
1003 const M5: &Rates = &Rates::AARCH64_APPLE;
1006
1007 fn entries(rt: &Rates) -> [(&'static str, Rate); 16] {
1009 [
1010 ("gemm_mac", rt.gemm_mac),
1011 ("gemm_single", rt.gemm_single),
1012 ("block_apply_mac", rt.block_apply_mac),
1013 ("lu_full_pivot", rt.lu_full_pivot),
1014 ("bunch_kaufman", rt.bunch_kaufman),
1015 ("qr_col_pivot", rt.qr_col_pivot),
1016 ("bidiag", rt.bidiag),
1017 ("hessenberg", rt.hessenberg),
1018 ("tridiag", rt.tridiag),
1019 ("matvec", rt.matvec),
1020 ("bidiag_svd_dc", rt.bidiag_svd_dc),
1021 ("tridiag_evd_dc", rt.tridiag_evd_dc),
1022 ("triangular_inverse", rt.triangular_inverse),
1023 ("sparse_matmul", rt.sparse_matmul),
1024 ("householder_rank_one", rt.householder_rank_one),
1025 ("lu_swap", rt.lu_swap),
1026 ]
1027 }
1028
1029 #[test]
1030 fn every_class_resolves_to_its_table() {
1031 assert_eq!(Rates::for_class(HostClass::X86_64Avx512), &Rates::X86_64_AVX512);
1032 assert_eq!(Rates::for_class(HostClass::Aarch64Apple), &Rates::AARCH64_APPLE);
1033 for c in [HostClass::X86_64Avx2, HostClass::X86_64Scalar, HostClass::Aarch64, HostClass::Generic] {
1034 assert_eq!(Rates::for_class(c), &Rates::GENERIC, "{}", c.name());
1035 }
1036 assert_eq!(Rates::X86_64_AVX512.gemm_mac, Rate { real: 0.0641, cplx: 0.2565 });
1038 assert_eq!(Rates::X86_64_AVX512.gemm_single, Rate { real: 0.0320, cplx: 0.1283 });
1039 assert_eq!(Rates::AARCH64_APPLE.gemm_mac, Rate { real: 0.037, cplx: 0.165 });
1040 assert_eq!(Rates::AARCH64_APPLE.bidiag_svd_dc, Rate { real: 53.0, cplx: 53.0 });
1041 let g = Rates::X86_64_AVX512;
1043 assert!((g.gemm_single.real / (g.gemm_mac.real / 2.0) - 1.0).abs() < 0.01);
1044 assert!((g.gemm_single.cplx / (g.gemm_mac.cplx / 2.0) - 1.0).abs() < 0.01);
1045 }
1046
1047 #[test]
1048 fn generic_is_the_per_entry_minimum_of_the_measured_rows() {
1049 let (a, b, g) = (entries(&Rates::X86_64_AVX512), entries(&Rates::AARCH64_APPLE), entries(&Rates::GENERIC));
1050 for i in 0..a.len() {
1051 let (name, ga) = (a[i].0, a[i].1);
1052 assert_eq!(g[i].1.real, ga.real.min(b[i].1.real), "{name} real");
1053 assert_eq!(g[i].1.cplx, ga.cplx.min(b[i].1.cplx), "{name} cplx");
1054 assert!(g[i].1.real <= ga.real && g[i].1.real <= b[i].1.real, "{name}");
1055 }
1056 }
1057
1058 #[test]
1059 fn rows_are_positive() {
1060 for c in HostClass::ALL {
1061 for (name, r) in entries(Rates::for_class(c)) {
1062 assert!(r.real > 0.0 && r.cplx > 0.0, "{} {name}", c.name());
1063 }
1064 }
1065 }
1066
1067 #[test]
1068 fn detection_is_stable_and_the_host_row_is_its_class() {
1069 let c = HostClass::detect();
1070 for _ in 0..4 {
1071 assert_eq!(HostClass::detect(), c);
1072 }
1073 assert_eq!(HostClass::detect_uncached(), c);
1074 assert_eq!(rates(), Rates::for_class(c));
1075 assert_eq!(rate::bidiag(), rates().bidiag);
1076 #[cfg(all(target_arch = "x86_64", feature = "std"))]
1077 if std::arch::is_x86_feature_detected!("avx512f") {
1078 assert_eq!(c, HostClass::X86_64Avx512);
1079 }
1080 #[cfg(all(target_arch = "aarch64", target_os = "macos"))]
1081 assert_eq!(c, HostClass::Aarch64Apple);
1082 }
1083
1084 #[test]
1086 fn host_class_agrees_with_threadteam() {
1087 let tt = |c: HostClass| match c {
1088 HostClass::X86_64Avx512 => threadteam::HostClass::X86_64Avx512,
1089 HostClass::X86_64Avx2 => threadteam::HostClass::X86_64Avx2,
1090 HostClass::X86_64Scalar => threadteam::HostClass::X86_64Scalar,
1091 HostClass::Aarch64Apple => threadteam::HostClass::Aarch64Apple,
1092 HostClass::Aarch64 => threadteam::HostClass::Aarch64,
1093 HostClass::Generic => threadteam::HostClass::Generic,
1094 };
1095 assert_eq!(tt(HostClass::detect()), threadteam::HostClass::detect());
1096 for c in HostClass::ALL {
1097 assert_eq!(c.name(), tt(c).name());
1098 }
1099 }
1100
1101 #[test]
1105 fn the_team_cost_is_threadteams_table_at_every_width() {
1106 use threadteam::{DispatchCosts, DispatchKind, HostClass, WidthCosts};
1107 let us = Duration::from_micros;
1108 let rows = [
1109 DispatchCosts::for_class(HostClass::Generic),
1110 DispatchCosts::for_class(HostClass::X86_64Avx512),
1111 DispatchCosts::for_class(HostClass::X86_64Avx2),
1112 DispatchCosts::for_class(HostClass::Aarch64Apple),
1113 DispatchCosts::for_class(HostClass::Generic).with(DispatchKind::Team, WidthCosts::flat(us(1))),
1114 ];
1115 for costs in rows {
1116 let s = ParShape::Team(nz(128), costs);
1117 for w in 2..=128 {
1118 assert_eq!(s.dispatch_cost(nz(w)), costs.cost(DispatchKind::Team, nz(w)), "{costs:?} width {w}");
1119 }
1120 }
1121 let flat = ParShape::Team(nz(128), rows[4]);
1124 for (w, ns) in [(8, 1333), (12, 2000), (16, 2666), (32, 5333), (64, 10666), (128, 21333)] {
1125 assert_eq!(flat.dispatch_cost(nz(w)), Duration::from_nanos(ns), "flat width {w}");
1126 }
1127 let avx512 = ParShape::Team(nz(128), rows[1]);
1128 assert_eq!(avx512.dispatch_cost(nz(8)), Duration::from_nanos(980));
1129 assert_eq!(avx512.dispatch_cost(nz(64)), Duration::from_nanos(980 * 8));
1130 let zen2 = ParShape::Team(nz(128), rows[2]);
1131 assert_eq!(zen2.dispatch_cost(nz(64)), Duration::from_nanos(9333));
1132 assert_eq!(zen2.dispatch_cost(nz(128)), Duration::from_nanos(18666));
1133 let at_8 = flat.dispatch_cost(nz(8)) * 3;
1135 assert!(flat.pays(at_8, 1, nz(8)));
1136 assert!(!flat.pays(at_8, 1, nz(64)));
1137 }
1138
1139 #[test]
1140 fn fixed_and_never_keep_their_meaning() {
1141 let r = rate::lu_full_pivot();
1142 assert_eq!(ParThreshold::Fixed(100).region_width::<f64>(&crate::Par::Seq, 100, r, 1), 1);
1144 crate::team::with_team(6, |par| {
1145 assert_eq!(ParThreshold::Never.region_width::<f64>(&par, usize::MAX, r, 1), 1);
1146 assert_eq!(ParThreshold::Fixed(100).region_width::<f64>(&par, 99, r, 1), 1);
1147 assert_eq!(ParThreshold::Fixed(100).region_width::<f64>(&par, 100, r, 1), 6);
1148 });
1149 assert_eq!(ParThreshold::from(0), ParThreshold::Derived);
1151 assert_eq!(ParThreshold::from(usize::MAX), ParThreshold::Never);
1152 assert_eq!(ParThreshold::from(7), ParThreshold::Fixed(7));
1153 assert_eq!(usize::from(ParThreshold::Fixed(0)), 1);
1154 }
1155
1156 #[test]
1157 fn derived_gates_never_pay_on_a_sequential_shape() {
1158 let seq = crate::Par::Seq;
1159 for (_, r, sections) in gates(rates()) {
1160 assert_eq!(ParThreshold::Derived.region_width::<f64>(&seq, usize::MAX / 2, r, sections), 1);
1161 assert_eq!(ParThreshold::Derived.region_width::<c64>(&seq, usize::MAX / 2, r, sections), 1);
1162 assert!(!ParShape::Seq.pays_full(serial_time(usize::MAX / 2, r.of::<f64>()), sections));
1163 }
1164 }
1165
1166 #[test]
1167 fn derived_floors_on_a_team_follow_the_rule() {
1168 use threadteam::DispatchCosts;
1169 let team = ParShape::Team(nz(6), DispatchCosts::for_class(threadteam::HostClass::Aarch64Apple));
1170 let d = team.dispatch_cost(nz(6)).as_secs_f64() * 1e9;
1171 for (name, r, sections) in gates(M5) {
1172 for (ty, ns) in [("f64", r.real), ("c64", r.cplx)] {
1173 let floor = if ty == "f64" {
1174 floor_units::<f64>(team, r, sections)
1175 } else {
1176 floor_units::<c64>(team, r, sections)
1177 };
1178 let want = PAYS * f64::from(sections) * d / ns;
1180 assert!(
1181 (floor as f64 - want).abs() <= 4.0 + want * 2e-3,
1182 "{name} {ty}: floor {floor}, rule {want}"
1183 );
1184 }
1185 }
1186 let pays = |u: usize, r: Rate, f: u32| {
1188 team.pays_full(serial_time(u, r.of::<f64>()), f)
1189 };
1190 assert!(!pays(100 * 100, M5.lu_full_pivot, 1));
1192 assert!(pays(200 * 200, M5.lu_full_pivot, 1));
1193 assert!(floor_units::<f64>(team, M5.bidiag, 2) >= 2 * floor_units::<f64>(team, M5.bidiag, 1) - 2);
1195 assert!(pays(256 * 256, M5.matvec, 1));
1197 assert!(!pays(64 * 64, M5.matvec, 1));
1198 let lu = floor_units::<f64>(team, M5.lu_full_pivot, 1);
1201 let bidiag = floor_units::<f64>(team, M5.bidiag, 1);
1202 assert!(lu > bidiag);
1203 }
1204
1205 #[test]
1206 fn par_width_reads_the_degree_of_the_token() {
1207 let seq = crate::Par::Seq;
1208 assert_eq!(par_width::<f64>(&seq, usize::MAX / 2, rate::matvec(), 1), 1);
1209 crate::team::with_team(6, |par| {
1210 assert_eq!(par_width::<f64>(&par, 10_000_000, rate::matvec(), 1), 6);
1211 assert_eq!(par_width::<f64>(&par, 1_000, rate::matvec(), 1), 1);
1212 assert_eq!(par_width_upto::<f64>(&par, 10_000_000, rate::matvec(), 1, 3), 3);
1214 });
1215 crate::team::with_team(1, |par| {
1216 assert_eq!(par_width::<f64>(&par, 10_000_000, rate::matvec(), 1), 1);
1217 });
1218 }
1219
1220 #[test]
1221 fn rates_are_per_dtype() {
1222 let r = rate::bidiag();
1223 assert_eq!(r.of::<f64>(), r.real);
1224 assert_eq!(r.of::<c64>(), r.cplx);
1225 assert_eq!(r.of::<f32>(), r.real / 2.0);
1226 assert_eq!(r.of::<crate::c32>(), r.cplx / 2.0);
1227 assert_eq!(rate::gemm_mac().of::<f64>(), ns_per_mac::f64());
1228 assert_eq!(rate::gemm_mac().of::<c64>(), ns_per_mac::c64());
1229 }
1230
1231 #[cfg(feature = "rayon")]
1235 mod forced_parallel {
1236 use super::*;
1237 use crate::dyn_stack::{MemBuffer, MemStack};
1238 use crate::internal_prelude::*;
1239 use crate::linalg::cholesky::lblt::factor as bk;
1240 use crate::linalg::evd::{hessenberg, tridiag};
1241 use crate::linalg::lu::full_pivoting::factor as fplu;
1242 use crate::linalg::qr::col_pivoting::factor as cpqr;
1243 use crate::linalg::svd::bidiag;
1244 use crate::{Auto, Mat, MatRef, Par};
1245
1246 const N: usize = 96;
1247
1248 fn fill<T: crate::ComplexField>(n: usize, seed: usize) -> Mat<T> {
1249 Mat::from_fn(n, n, |i, j| {
1250 let h = (i * 2654435761usize)
1251 .wrapping_add(j * 40503)
1252 .wrapping_add(seed * 977)
1253 % 1009;
1254 from_f64::<T>(h as f64 / 504.0 - 1.0)
1255 })
1256 }
1257 fn close<T: crate::ComplexField<Real = f64>>(
1258 what: &str,
1259 a: MatRef<'_, T>,
1260 b: MatRef<'_, T>,
1261 ) {
1262 let mut err = 0.0f64;
1263 for j in 0..a.ncols() {
1264 for i in 0..a.nrows() {
1265 err = err.max((a[(i, j)].clone() - b[(i, j)].clone()).abs());
1266 }
1267 }
1268 core::assert!(err <= 1e-9, "{what}: error {err}");
1269 }
1270
1271 fn check<T: crate::ComplexField<Real = f64>>(compare_complex_factors: bool) {
1272 let a0 = fill::<T>(N, 1);
1273 let sym = &a0 + a0.adjoint();
1274 let lu = |par: Par| {
1276 let mut a = a0.clone();
1277 let mut p = vec![0usize; 4 * N];
1278 let (r, rest) = p.split_at_mut(N);
1279 let (ri, rest) = rest.split_at_mut(N);
1280 let (c, ci) = rest.split_at_mut(N);
1281 let mut params: fplu::FullPivLuParams = Auto::<T>::auto();
1282 params.par_threshold = ParThreshold::Fixed(1);
1283 let mut mem = MemBuffer::new(fplu::lu_in_place_scratch::<usize, T>(
1284 N,
1285 N,
1286 par.shape(),
1287 params.into(),
1288 ));
1289 fplu::lu_in_place(
1290 a.as_mut(), r, ri, c, ci, par, MemStack::new(&mut mem), params.into(),
1291 );
1292 a
1293 };
1294 crate::team::with_team(4, |par| close("lu", lu(Par::Seq).as_ref(), lu(par).as_ref()));
1295 let bunch = |par: Par| {
1296 let mut a = sym.clone();
1297 let mut perm = vec![0usize; N];
1298 let mut perm_inv = vec![0usize; N];
1299 let mut sub = Mat::<T>::zeros(N, 1);
1300 let mut params: bk::LbltParams = Auto::<T>::auto();
1301 params.pivoting = bk::PivotingStrategy::Full;
1302 params.par_threshold = ParThreshold::Fixed(1);
1303 let mut mem = MemBuffer::new(bk::cholesky_in_place_scratch::<usize, T>(
1304 N,
1305 par.shape(),
1306 params.into(),
1307 ));
1308 bk::cholesky_in_place(
1309 a.as_mut(),
1310 sub.as_mut().col_mut(0).as_diagonal_mut(),
1311 &mut perm,
1312 &mut perm_inv,
1313 par,
1314 MemStack::new(&mut mem),
1315 params.into(),
1316 );
1317 a
1318 };
1319 crate::team::with_team(4, |par| close("bunch", bunch(Par::Seq).as_ref(), bunch(par).as_ref()));
1320 let qr = |par: Par| {
1321 let mut a = a0.clone();
1322 let mut h = Mat::<T>::zeros(16, N);
1323 let mut p = vec![0usize; N];
1324 let mut pi = vec![0usize; N];
1325 let mut params: cpqr::ColPivQrParams = Auto::<T>::auto();
1326 params.par_threshold = ParThreshold::Fixed(1);
1327 let mut mem = MemBuffer::new(cpqr::qr_in_place_scratch::<usize, T>(
1328 N,
1329 N,
1330 16,
1331 par.shape(),
1332 params.into(),
1333 ));
1334 cpqr::qr_in_place(
1335 a.as_mut(), h.as_mut(), &mut p, &mut pi, par, MemStack::new(&mut mem), params.into(),
1336 );
1337 a
1338 };
1339 crate::team::with_team(4, |par| close("qr", qr(Par::Seq).as_ref(), qr(par).as_ref()));
1340 let bid = |par: Par| {
1341 let mut a = a0.clone();
1342 let mut hl = Mat::<T>::zeros(8, N);
1343 let mut hr = Mat::<T>::zeros(8, N - 1);
1344 let mut params: bidiag::BidiagParams = Auto::<T>::auto();
1345 params.par_threshold = ParThreshold::Fixed(1);
1346 let mut mem = MemBuffer::new(bidiag::bidiag_in_place_scratch::<T>(
1347 N,
1348 N,
1349 par.shape(),
1350 params.into(),
1351 ));
1352 bidiag::bidiag_in_place(
1353 a.as_mut(), hl.as_mut(), hr.as_mut(), par, MemStack::new(&mut mem), params.into(),
1354 );
1355 for j in 0..N {
1360 for i in 0..N {
1361 if i > j || j > i + 1 {
1362 a[(i, j)] = T::zero();
1363 }
1364 }
1365 }
1366 a
1367 };
1368 if compare_complex_factors {
1375 crate::team::with_team(4, |par| {
1376 close("bid", bid(Par::Seq).as_ref(), bid(par).as_ref())
1377 });
1378 } else {
1379 let seq = bid(Par::Seq);
1383 let team = crate::team::with_team(4, |par| bid(par));
1384 core::assert!(seq.as_ref().is_all_finite(), "bid sequential");
1385 core::assert!(team.as_ref().is_all_finite(), "bid team");
1386 }
1387 let hess = |par: Par| {
1388 let mut a = a0.clone();
1389 let mut h = Mat::<T>::zeros(8, N - 1);
1390 let mut params: hessenberg::HessenbergParams = Auto::<T>::auto();
1391 params.par_threshold = ParThreshold::Fixed(1);
1392 params.blocking_threshold = usize::MAX;
1393 let mut mem = MemBuffer::new(hessenberg::hessenberg_in_place_scratch::<T>(
1394 N,
1395 8,
1396 par.shape(),
1397 params.into(),
1398 ));
1399 hessenberg::hessenberg_in_place(
1400 a.as_mut(), h.as_mut(), par, MemStack::new(&mut mem), params.into(),
1401 );
1402 a
1403 };
1404 if compare_complex_factors {
1409 crate::team::with_team(4, |par| {
1410 close("hess", hess(Par::Seq).as_ref(), hess(par).as_ref())
1411 });
1412 } else {
1413 let seq = hess(Par::Seq);
1414 let team = crate::team::with_team(4, |par| hess(par));
1415 core::assert!(seq.as_ref().is_all_finite(), "hess sequential");
1416 core::assert!(team.as_ref().is_all_finite(), "hess team");
1417 }
1418 let tri = |par: Par| {
1419 let mut a = sym.clone();
1420 let mut h = Mat::<T>::zeros(8, N - 1);
1421 let mut params: tridiag::TridiagParams = Auto::<T>::auto();
1422 params.par_threshold = ParThreshold::Fixed(1);
1423 let mut mem = MemBuffer::new(tridiag::tridiag_in_place_scratch::<T>(
1424 N,
1425 par.shape(),
1426 params.into(),
1427 ));
1428 tridiag::tridiag_in_place(
1429 a.as_mut(), h.as_mut(), par, MemStack::new(&mut mem), params.into(),
1430 );
1431 a
1432 };
1433 crate::team::with_team(4, |par| close("tri", tri(Par::Seq).as_ref(), tri(par).as_ref()));
1434 }
1435
1436 #[test]
1437 fn forced_parallel_steps_equal_serial_f64() {
1438 check::<f64>(true);
1439 }
1440
1441 #[test]
1442 fn forced_parallel_steps_equal_serial_c64() {
1443 check::<c64>(false);
1444 }
1445 }
1446
1447 #[test]
1449 fn halves_join_gates_follow_the_rule() {
1450 let seq = crate::Par::Seq;
1451 for r in [rate::bidiag_svd_dc(), rate::tridiag_evd_dc(), rate::triangular_inverse()] {
1452 assert!(!halves_pays::<f64>(&seq, 1 << 20, r, 2));
1453 }
1454 crate::team::with_team(6, |par| {
1455 for r in [rate::bidiag_svd_dc(), rate::tridiag_evd_dc(), rate::triangular_inverse()] {
1456 let power = if r == rate::triangular_inverse() { 3 } else { 2 };
1458 let first = (2..4096).find(|&n| halves_pays::<f64>(&par, n, r, power)).expect("never pays");
1459 assert!(first > 1, "{first}");
1462 assert!(!halves_pays::<f64>(&par, first - 1, r, power));
1463 assert!(halves_pays::<f64>(&par, first, r, power));
1464 assert!(halves_pays::<f64>(&par, 2 * first, r, power));
1465 let h = first / 2;
1467 let work_ns = 2.0 * (h as f64).powi(power as i32) * r.real;
1468 let want = PAYS * par.shape().dispatch_cost(nz(2)).as_secs_f64() * 1e9;
1469 assert!(work_ns >= want * 0.999, "{work_ns} {want}");
1470 }
1471 });
1472 }
1473
1474 #[test]
1481 fn block_apply_split_starts_where_the_estimate_pays() {
1482 use threadteam::DispatchCosts;
1483 let team = ParShape::Team(nz(6), DispatchCosts::for_class(threadteam::HostClass::Aarch64Apple));
1484 let pays = |m: usize| {
1485 let (n, k) = (8, m - 8);
1486 pays_at(block_apply_estimate_with::<f64>(M5, m, n, k), 1, team.dispatch_cost(nz(6)))
1487 };
1488 assert!(!pays(72));
1489 assert!(pays(80));
1490 let gemm_only = |m: usize| {
1492 pays_at(serial_time(2 * m * 8 * (m - 8), M5.gemm_mac.real), 1, team.dispatch_cost(nz(6)))
1493 };
1494 assert!(!gemm_only(88) && gemm_only(96));
1495 let c = block_apply_estimate_with::<c64>(M5, 100, 8, 92).as_secs_f64() * 1e9;
1497 let want = 2.0 * 100.0 * 8.0 * 92.0 * M5.block_apply_mac.cplx + 100.0 * 92.0 * M5.householder_rank_one.cplx;
1498 assert!((c - want).abs() < 2.0, "{c} {want}");
1499 }
1500
1501 #[test]
1502 fn join_gates_on_a_team_follow_the_width_two_cost() {
1503 use threadteam::DispatchCosts;
1504 let team = ParShape::Team(nz(6), DispatchCosts::for_class(threadteam::HostClass::Aarch64Apple));
1505 for (name, r, sections) in gates(M5) {
1506 assert!(floor_units::<f64>(team, r, sections) > 0, "{name}");
1507 }
1508 let d2 = team.dispatch_cost(nz(2));
1510 assert_eq!(d2, Duration::from_micros(3));
1511 let par = |ns: f64| pays_at(Duration::from_secs_f64(ns * 1e-9), 1, d2);
1512 assert!(par(9000.0) && !par(8900.0));
1513 for (r, power, yes, no) in [
1517 (M5.bidiag_svd_dc, 2u32, 128usize, 8usize),
1518 (M5.tridiag_evd_dc, 2, 64, 8),
1519 (M5.triangular_inverse, 3, 256, 64),
1520 ] {
1521 let pays = |n: usize| {
1522 let h = n / 2;
1523 pays_at(
1524 serial_time(2 * h.pow(power), r.real),
1525 1,
1526 d2,
1527 )
1528 };
1529 assert!(pays(yes));
1530 assert!(!pays(no));
1531 }
1532 assert!(team.pays_full(serial_time(1_000_000, M5.sparse_matmul.of::<f64>()), 1));
1534 assert!(!team.pays_full(serial_time(1_000, M5.sparse_matmul.of::<f64>()), 1));
1535 }
1536}
1537
1538#[cfg(feature = "par-trace")]
1539#[doc(hidden)]
1540pub fn trace_pub(
1541 units: usize,
1542 sections: u32,
1543 estimate: Duration,
1544 dispatch: Duration,
1545 width: usize,
1546 r: bool,
1547 loc: &'static core::panic::Location<'static>,
1548) {
1549 trace::record(loc, units, sections, estimate.as_secs_f64() * 1e9, dispatch.as_secs_f64() * 1e9, width, r);
1550}
1551
1552#[cfg(all(test, feature = "sparse-linalg"))]
1553mod join_results {
1554 use crate::dyn_stack::{MemBuffer, MemStack};
1557 use crate::linalg::evd::{self, ComputeEigenvectors};
1558 use crate::linalg::svd;
1559 use crate::linalg::triangular_inverse::{invert_lower_triangular, invert_unit_lower_triangular};
1560 use crate::{Accum, Mat, Par, Spec};
1561
1562 struct Rng(u64);
1563 impl Rng {
1564 fn next(&mut self) -> f64 {
1565 self.0 ^= self.0 << 13;
1566 self.0 ^= self.0 >> 7;
1567 self.0 ^= self.0 << 17;
1568 (self.0 >> 11) as f64 / (1u64 << 53) as f64 - 0.5
1569 }
1570 }
1571
1572 fn close(a: &[f64], b: &[f64], tol: f64) {
1573 let scale = a.iter().fold(1.0f64, |m, x| m.max(x.abs()));
1574 for (x, y) in a.iter().zip(b) {
1575 assert!((x - y).abs() <= tol * scale, "{x} {y}");
1576 }
1577 }
1578
1579 #[test]
1580 fn triangular_inverse_agrees_with_serial() {
1581 let n = 512;
1582 let mut r = Rng(3);
1583 let mut src = Mat::<f64>::from_fn(n, n, |_, _| r.next() * 0.1);
1584 for i in 0..n {
1585 src[(i, i)] = 2.0 + src[(i, i)].abs();
1586 }
1587 for unit in [false, true] {
1588 let mut a = Mat::<f64>::zeros(n, n);
1589 let mut b = Mat::<f64>::zeros(n, n);
1590 if unit {
1591 invert_unit_lower_triangular(a.as_mut(), src.as_ref(), Par::Seq);
1592 } else {
1593 invert_lower_triangular(a.as_mut(), src.as_ref(), Par::Seq);
1594 }
1595 crate::team::with_team(4, |par| {
1596 if unit {
1597 invert_unit_lower_triangular(b.as_mut(), src.as_ref(), par);
1598 } else {
1599 invert_lower_triangular(b.as_mut(), src.as_ref(), par);
1600 }
1601 });
1602 for j in 0..n {
1603 close(
1604 &(0..n).map(|i| a[(i, j)]).collect::<Vec<_>>(),
1605 &(0..n).map(|i| b[(i, j)]).collect::<Vec<_>>(),
1606 1e-9,
1607 );
1608 }
1609 }
1610 }
1611
1612 #[test]
1613 fn bidiagonal_svd_agrees_with_serial() {
1614 let n = 512;
1615 let mut r = Rng(5);
1616 let d0: Vec<f64> = (0..n).map(|_| r.next() + 1.0).collect();
1617 let e: Vec<f64> = (0..n - 1).map(|_| r.next()).collect();
1618 let mut out = vec![];
1619 let run = |par: Par, out: &mut Vec<Vec<f64>>| {
1620 let mut d = d0.clone();
1621 let mut u = Mat::<f64>::zeros(n, n);
1622 let mut v = Mat::<f64>::zeros(n, n);
1623 let mut mem = MemBuffer::new(svd::bidiagonal_svd_scratch::<f64>(
1624 n,
1625 par.shape(),
1626 Spec::default(),
1627 ));
1628 svd::bidiagonal_svd(
1629 &mut d,
1630 &e,
1631 true,
1632 u.as_mut(),
1633 v.as_mut(),
1634 par,
1635 MemStack::new(&mut mem),
1636 Spec::default(),
1637 )
1638 .unwrap();
1639 out.push(d);
1640 };
1641 run(Par::Seq, &mut out);
1642 crate::team::with_team(4, |par| run(par, &mut out));
1643 close(&out[0], &out[1], 1e-10);
1644 }
1645
1646 #[test]
1647 fn tridiagonal_evd_agrees_with_serial() {
1648 let n = 512;
1649 let mut r = Rng(6);
1650 let d0: Vec<f64> = (0..n).map(|_| r.next()).collect();
1651 let mut e0: Vec<f64> = (0..n - 1).map(|_| r.next()).collect();
1652 e0.push(0.0);
1653 let mut out = vec![];
1654 let run = |par: Par, out: &mut Vec<Vec<f64>>| {
1655 let mut u = Mat::<f64>::zeros(n, n);
1656 let mut s = Mat::<f64>::zeros(n, 1);
1657 let mut mem = MemBuffer::new(evd::tridiagonal_self_adjoint_evd_scratch::<f64>(
1658 n,
1659 ComputeEigenvectors::Yes,
1660 par.shape(),
1661 Spec::default(),
1662 ));
1663 let d = crate::col::ColRef::from_slice(&d0);
1664 let e = crate::col::ColRef::from_slice(&e0);
1665 evd::tridiagonal_self_adjoint_evd(
1666 d.as_diagonal(),
1667 e.as_diagonal(),
1668 s.as_mut().col_mut(0).as_diagonal_mut(),
1669 Some(u.as_mut()),
1670 par,
1671 MemStack::new(&mut mem),
1672 Spec::default(),
1673 )
1674 .unwrap();
1675 out.push((0..n).map(|i| s[(i, 0)]).collect::<Vec<_>>());
1676 };
1677 run(Par::Seq, &mut out);
1678 crate::team::with_team(4, |par| run(par, &mut out));
1679 close(&out[0], &out[1], 1e-10);
1680 }
1681
1682 #[test]
1683 fn sparse_matmul_agrees_with_serial() {
1684 use crate::sparse::SparseColMat;
1685 use crate::sparse::Triplet;
1686 use crate::sparse::linalg::matmul::sparse_sparse_matmul;
1687 let n = 6000;
1688 let mk = |seed: u64| {
1689 let mut r = Rng(seed);
1690 let mut t = vec![];
1691 for j in 0..n {
1692 for _ in 0..24 {
1693 let i = ((r.next() + 0.5) * n as f64) as usize % n;
1694 t.push(Triplet::new(i, j, r.next()));
1695 }
1696 }
1697 SparseColMat::<usize, f64>::try_new_from_triplets(n, n, &t).unwrap()
1698 };
1699 let (a, b) = (mk(1), mk(2));
1700 let s = sparse_sparse_matmul(a.as_ref(), b.as_ref(), 1.0, Par::Seq).unwrap();
1701 let p = crate::team::with_team(4, |par| {
1702 sparse_sparse_matmul(a.as_ref(), b.as_ref(), 1.0, par).unwrap()
1703 });
1704 close(s.val(), p.val(), 1e-12);
1705 let _ = Accum::Add;
1706 }
1707}