1use crate::ParShape;
11use core::num::NonZeroUsize;
12use core::time::Duration;
13
14#[cfg(feature = "threadteam")]
15pub use threadteam::{PAYS, pays_at, width_for_by};
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)]
44pub enum HostClass {
45 X86_64Avx512,
47 X86_64Avx2,
49 X86_64Scalar,
51 Aarch64Apple,
53 Aarch64,
55 Generic,
57}
58
59impl HostClass {
60 pub const ALL: [HostClass; 6] = [
62 HostClass::X86_64Avx512,
63 HostClass::X86_64Avx2,
64 HostClass::X86_64Scalar,
65 HostClass::Aarch64Apple,
66 HostClass::Aarch64,
67 HostClass::Generic,
68 ];
69
70 pub fn detect() -> Self {
73 use core::sync::atomic::{AtomicU8, Ordering};
74 static CLASS: AtomicU8 = AtomicU8::new(0);
75 let c = CLASS.load(Ordering::Relaxed);
76 if c != 0 {
77 return Self::ALL[c as usize - 1];
78 }
79 let class = Self::detect_uncached();
80 let idx = Self::ALL.iter().position(|&k| k == class).unwrap();
81 CLASS.store(idx as u8 + 1, Ordering::Relaxed);
82 class
83 }
84
85 #[cfg(all(target_arch = "x86_64", feature = "std"))]
86 fn detect_uncached() -> Self {
87 if std::arch::is_x86_feature_detected!("avx512f") {
88 HostClass::X86_64Avx512
89 } else if std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("fma") {
90 HostClass::X86_64Avx2
91 } else {
92 HostClass::X86_64Scalar
93 }
94 }
95
96 #[cfg(all(target_arch = "x86_64", not(feature = "std")))]
97 fn detect_uncached() -> Self {
98 if cfg!(target_feature = "avx512f") {
99 HostClass::X86_64Avx512
100 } else if cfg!(all(target_feature = "avx2", target_feature = "fma")) {
101 HostClass::X86_64Avx2
102 } else {
103 HostClass::X86_64Scalar
104 }
105 }
106
107 #[cfg(all(target_arch = "aarch64", target_os = "macos"))]
108 fn detect_uncached() -> Self {
109 HostClass::Aarch64Apple
110 }
111
112 #[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
113 fn detect_uncached() -> Self {
114 HostClass::Aarch64
115 }
116
117 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
118 fn detect_uncached() -> Self {
119 HostClass::Generic
120 }
121
122 pub const fn name(self) -> &'static str {
124 match self {
125 HostClass::X86_64Avx512 => "x86_64-avx512",
126 HostClass::X86_64Avx2 => "x86_64-avx2",
127 HostClass::X86_64Scalar => "x86_64-scalar",
128 HostClass::Aarch64Apple => "aarch64-apple",
129 HostClass::Aarch64 => "aarch64",
130 HostClass::Generic => "generic",
131 }
132 }
133}
134
135#[cfg(any(all(feature = "rayon", not(feature = "threadteam")), test))]
147#[cfg_attr(feature = "threadteam", allow(dead_code))]
148pub(crate) fn outside_cost_of(class: HostClass, w: NonZeroUsize) -> Duration {
149 width_model(outside_anchors(class), w)
150}
151
152#[cfg(any(all(feature = "rayon", not(feature = "threadteam")), test))]
153#[cfg_attr(feature = "threadteam", allow(dead_code))]
154const fn outside_anchors(class: HostClass) -> [u64; 3] {
155 match class {
156 HostClass::X86_64Avx512 => [9_700, 10_100, 8_700],
157 HostClass::Aarch64Apple => [14_000, 18_000, 19_000],
158 _ => [40_000; 3],
159 }
160}
161
162#[cfg(any(all(feature = "rayon", not(feature = "threadteam")), test))]
177#[cfg_attr(feature = "threadteam", allow(dead_code))]
178pub(crate) fn in_pool_cost_of(class: HostClass, w: NonZeroUsize) -> Duration {
179 width_model(in_pool_anchors(class), w)
180}
181
182#[cfg(any(all(feature = "rayon", not(feature = "threadteam")), test))]
183#[cfg_attr(feature = "threadteam", allow(dead_code))]
184const fn in_pool_anchors(class: HostClass) -> [u64; 3] {
185 match class {
186 HostClass::X86_64Avx512 => [3_000, 3_600, 3_500],
187 HostClass::Aarch64Apple => [6_000, 11_000, 5_500],
188 _ => [15_000; 3],
189 }
190}
191
192#[cfg(any(all(feature = "rayon", not(feature = "threadteam")), test))]
195#[cfg_attr(feature = "threadteam", allow(dead_code))]
196fn width_model(anchors: [u64; 3], w: NonZeroUsize) -> Duration {
197 let w = w.get() as i64;
198 let ns = if w <= 2 {
199 anchors[0] as i64
200 } else if w >= 6 {
201 anchors[2] as i64
202 } else {
203 let i = (w > 4) as usize;
204 let (w0, w1) = (2 + 2 * i as i64, 4 + 2 * i as i64);
205 let (c0, c1) = (anchors[i] as i64, anchors[i + 1] as i64);
206 c0 + (c1 - c0) * (w - w0) / (w1 - w0)
207 };
208 Duration::from_nanos(ns as u64)
209}
210
211#[cfg(feature = "rayon")]
216#[inline]
217pub(crate) fn caller_in_pool() -> bool {
218 rayon::current_thread_index().is_some()
219}
220
221#[cfg(all(feature = "rayon", not(feature = "threadteam")))]
223pub(crate) fn rayon_cost(w: NonZeroUsize) -> Duration {
224 if caller_in_pool() {
225 in_pool_cost_of(HostClass::detect(), w)
226 } else {
227 outside_cost_of(HostClass::detect(), w)
228 }
229}
230
231#[cfg(all(feature = "rayon", not(feature = "threadteam")))]
235pub(crate) fn min_rayon_cost() -> Duration {
236 let class = HostClass::detect();
237 let anchors = if caller_in_pool() { in_pool_anchors(class) } else { outside_anchors(class) };
238 Duration::from_nanos(anchors[0].min(anchors[1]).min(anchors[2]))
239}
240
241pub mod ns_per_mac {
244 use super::rates;
245 pub fn f64() -> f64 {
247 rates().gemm_mac.real
248 }
249 pub fn c64() -> f64 {
251 rates().gemm_mac.cplx
252 }
253 pub fn f32() -> f64 {
255 rates().gemm_single.real
256 }
257 pub fn c32() -> f64 {
259 rates().gemm_single.cplx
260 }
261}
262
263pub const NS_PER_SWAP: f64 = 0.55;
269
270#[derive(Copy, Clone, Debug, PartialEq, Eq, Default)]
278pub enum ParThreshold {
279 #[default]
281 Derived,
282 Fixed(usize),
287 Never,
289}
290
291impl From<usize> for ParThreshold {
295 fn from(value: usize) -> Self {
296 match value {
297 0 => ParThreshold::Derived,
298 usize::MAX => ParThreshold::Never,
299 n => ParThreshold::Fixed(n),
300 }
301 }
302}
303
304impl From<ParThreshold> for usize {
307 fn from(value: ParThreshold) -> Self {
308 match value {
309 ParThreshold::Derived => 0,
310 ParThreshold::Fixed(n) => n.max(1),
311 ParThreshold::Never => usize::MAX,
312 }
313 }
314}
315
316impl ParThreshold {
317 #[cfg_attr(feature = "par-trace", track_caller)]
323 pub(crate) fn region_pays<T: crate::ComplexField>(
324 self,
325 shape: ParShape,
326 units: usize,
327 rate: Rate,
328 sections: u32,
329 ) -> bool {
330 if shape.degree() < 2 && !matches!(self, Self::Fixed(_)) {
333 return false;
334 }
335 let r = self.region_pays_impl::<T>(shape, units, rate, sections);
336 #[cfg(feature = "par-trace")]
337 trace_decision(shape, shape.degree(), units, rate.of::<T>(), sections, r, core::panic::Location::caller());
338 r
339 }
340
341 fn region_pays_impl<T: crate::ComplexField>(
342 self,
343 shape: ParShape,
344 units: usize,
345 rate: Rate,
346 sections: u32,
347 ) -> bool {
348 match self {
349 Self::Never => false,
350 Self::Fixed(floor) => units >= floor,
351 Self::Derived => {
352 shape.pays_full(serial_time(units, rate.of::<T>()), sections)
353 },
354 }
355 }
356}
357
358#[derive(Copy, Clone, Debug, PartialEq)]
364pub struct Rate {
365 pub real: f64,
367 pub cplx: f64,
369}
370
371impl Rate {
372 pub fn of<T: crate::ComplexField>(self) -> f64 {
374 if const { T::IS_NATIVE_F32 } {
375 self.real / 2.0
376 } else if const { T::IS_NATIVE_C32 } {
377 self.cplx / 2.0
378 } else if const { T::IS_REAL } {
379 self.real
380 } else {
381 self.cplx
382 }
383 }
384}
385
386#[derive(Copy, Clone, Debug, PartialEq)]
408pub struct Rates {
409 pub gemm_mac: Rate,
412 pub gemm_single: Rate,
414 pub block_apply_mac: Rate,
416 pub lu_full_pivot: Rate,
419 pub bunch_kaufman: Rate,
422 pub qr_col_pivot: Rate,
425 pub bidiag: Rate,
428 pub hessenberg: Rate,
431 pub tridiag: Rate,
434 pub matvec: Rate,
437 pub bidiag_svd_dc: Rate,
444 pub tridiag_evd_dc: Rate,
448 pub triangular_inverse: Rate,
453 pub sparse_matmul: Rate,
456 pub householder_rank_one: Rate,
458 pub lu_swap: Rate,
462}
463
464const fn fmin(a: f64, b: f64) -> f64 {
465 if b < a { b } else { a }
466}
467
468const fn rmin(a: Rate, b: Rate) -> Rate {
469 Rate { real: fmin(a.real, b.real), cplx: fmin(a.cplx, b.cplx) }
470}
471
472const fn r(real: f64, cplx: f64) -> Rate {
473 Rate { real, cplx }
474}
475
476impl Rates {
477 const fn min_of(a: &Rates, b: &Rates) -> Rates {
479 Rates {
480 gemm_mac: rmin(a.gemm_mac, b.gemm_mac),
481 gemm_single: rmin(a.gemm_single, b.gemm_single),
482 block_apply_mac: rmin(a.block_apply_mac, b.block_apply_mac),
483 lu_full_pivot: rmin(a.lu_full_pivot, b.lu_full_pivot),
484 bunch_kaufman: rmin(a.bunch_kaufman, b.bunch_kaufman),
485 qr_col_pivot: rmin(a.qr_col_pivot, b.qr_col_pivot),
486 bidiag: rmin(a.bidiag, b.bidiag),
487 hessenberg: rmin(a.hessenberg, b.hessenberg),
488 tridiag: rmin(a.tridiag, b.tridiag),
489 matvec: rmin(a.matvec, b.matvec),
490 bidiag_svd_dc: rmin(a.bidiag_svd_dc, b.bidiag_svd_dc),
491 tridiag_evd_dc: rmin(a.tridiag_evd_dc, b.tridiag_evd_dc),
492 triangular_inverse: rmin(a.triangular_inverse, b.triangular_inverse),
493 sparse_matmul: rmin(a.sparse_matmul, b.sparse_matmul),
494 householder_rank_one: rmin(a.householder_rank_one, b.householder_rank_one),
495 lu_swap: rmin(a.lu_swap, b.lu_swap),
496 }
497 }
498
499 pub const X86_64_AVX512: Rates = Rates {
501 gemm_mac: r(0.0641, 0.2565),
502 gemm_single: r(0.0320, 0.1283),
503 block_apply_mac: r(0.0676, 0.2685),
504 lu_full_pivot: r(0.4011, 0.8679),
505 bunch_kaufman: r(0.8295, 1.3904),
506 qr_col_pivot: r(0.3645, 1.0320),
507 bidiag: r(1.0427, 2.1997),
508 hessenberg: r(0.8576, 2.2212),
509 tridiag: r(0.7695, 1.9252),
510 matvec: r(0.1654, 0.4215),
512 bidiag_svd_dc: r(152.9, 152.9),
513 tridiag_evd_dc: r(97.4, 97.4),
514 triangular_inverse: r(0.01936, 0.07228),
515 sparse_matmul: r(4.1315, 5.2342),
517 householder_rank_one: r(0.4783, 1.0303),
518 lu_swap: r(0.98, 1.30),
522 };
523
524 pub const AARCH64_APPLE: Rates = Rates {
526 gemm_mac: r(0.037, 0.165),
527 gemm_single: r(0.037 / 2.0, 0.165 / 2.0),
528 block_apply_mac: r(0.037, 0.165),
529 lu_full_pivot: r(0.216, 0.6),
530 bunch_kaufman: r(0.238, 0.38),
531 qr_col_pivot: r(0.202, 0.713),
532 bidiag: r(0.48, 1.36),
533 hessenberg: r(0.451, 1.57),
534 tridiag: r(0.339, 1.16),
535 matvec: r(0.1, 0.34),
536 bidiag_svd_dc: r(53.0, 53.0),
537 tridiag_evd_dc: r(31.0, 31.0),
538 triangular_inverse: r(0.0096, 0.0344),
539 sparse_matmul: r(1.5, 1.8),
540 householder_rank_one: r(0.212, 0.725),
541 lu_swap: r(NS_PER_SWAP, NS_PER_SWAP),
542 };
543
544 pub const GENERIC: Rates = Rates::min_of(&Rates::X86_64_AVX512, &Rates::AARCH64_APPLE);
546
547 pub fn for_class(class: HostClass) -> &'static Rates {
549 static AVX512: Rates = Rates::X86_64_AVX512;
550 static APPLE: Rates = Rates::AARCH64_APPLE;
551 static GENERIC: Rates = Rates::GENERIC;
552 match class {
553 HostClass::X86_64Avx512 => &AVX512,
554 HostClass::Aarch64Apple => &APPLE,
555 HostClass::X86_64Avx2 | HostClass::X86_64Scalar | HostClass::Aarch64 | HostClass::Generic => &GENERIC,
556 }
557 }
558}
559
560pub fn rates() -> &'static Rates {
562 Rates::for_class(HostClass::detect())
563}
564
565pub mod rate {
568 use super::{Rate, rates};
569 macro_rules! families {
570 ($($name:ident),* $(,)?) => {$(
571 #[doc = core::concat!("`Rates::", core::stringify!($name), "` of the host class.")]
572 pub fn $name() -> Rate { rates().$name }
573 )*};
574 }
575 families!(
576 gemm_mac,
577 block_apply_mac,
578 lu_full_pivot,
579 bunch_kaufman,
580 qr_col_pivot,
581 bidiag,
582 hessenberg,
583 tridiag,
584 matvec,
585 bidiag_svd_dc,
586 tridiag_evd_dc,
587 triangular_inverse,
588 sparse_matmul,
589 householder_rank_one,
590 lu_swap,
591 );
592}
593
594pub(crate) fn gemm_ns_per_mac<T: crate::ComplexField>() -> f64 {
598 if const { T::IS_NATIVE_F32 } {
599 ns_per_mac::f32()
600 } else if const { T::IS_NATIVE_C32 } {
601 ns_per_mac::c32()
602 } else if const { T::IS_REAL } {
603 ns_per_mac::f64()
604 } else {
605 ns_per_mac::c64()
606 }
607}
608
609
610#[cfg(feature = "par-trace")]
614pub mod trace {
615 use std::sync::Mutex;
616 use std::sync::atomic::{AtomicBool, Ordering};
617 #[derive(Clone, Debug)]
619 pub struct Record {
620 pub site: &'static str,
622 pub dims: [usize; 3],
624 pub units: usize,
626 pub sections: u32,
628 pub estimate_ns: f64,
630 pub dispatch_ns: f64,
632 pub parallel: bool,
634 }
635 pub static N_TASKS_MODE: core::sync::atomic::AtomicUsize = core::sync::atomic::AtomicUsize::new(0);
638 static ON: AtomicBool = AtomicBool::new(false);
639 static LOG: Mutex<Vec<Record>> = Mutex::new(Vec::new());
640 std::thread_local! {
641 static DIMS: core::cell::Cell<(&'static str, [usize; 3])> = const { core::cell::Cell::new(("", [0; 3])) };
642 }
643 pub fn enable(on: bool) {
645 if on {
646 LOG.lock().unwrap().clear();
647 }
648 ON.store(on, Ordering::SeqCst);
649 }
650 pub fn take() -> Vec<Record> {
652 core::mem::take(&mut *LOG.lock().unwrap())
653 }
654 pub fn shape(label: &'static str, a: usize, b: usize, c: usize) {
656 DIMS.with(|d| d.set((label, [a, b, c])));
657 }
658 pub(crate) fn record(
659 loc: &'static core::panic::Location<'static>,
660 units: usize,
661 sections: u32,
662 estimate_ns: f64,
663 dispatch_ns: f64,
664 parallel: bool,
665 ) {
666 if !ON.load(Ordering::Relaxed) {
667 return;
668 }
669 let (label, dims) = DIMS.with(|d| d.replace(("", [0; 3])));
670 let site: &'static str = if label.is_empty() {
671 Box::leak(format!("{}:{}", loc.file().rsplit('/').next().unwrap(), loc.line()).into_boxed_str())
672 } else {
673 label
674 };
675 LOG.lock().unwrap().push(Record { site, dims, units, sections, estimate_ns, dispatch_ns, parallel });
676 }
677}
678
679pub(crate) fn par_pays_time(par: &crate::Par, estimate: Duration, sections: u32, width: usize) -> bool {
684 let width = width.min(par.degree());
685 let Some(width) = NonZeroUsize::new(width).filter(|w| w.get() >= 2) else {
686 return false;
687 };
688 match par {
689 crate::Par::Seq => false,
690 #[cfg(feature = "threadteam")]
691 crate::Par::Team(team) => team.pays(estimate, sections, width),
692 other => other.shape().pays(estimate, sections, width),
693 }
694}
695
696#[cfg(all(feature = "rayon", feature = "threadteam"))]
698fn rayon_kind() -> threadteam::DispatchKind {
699 if caller_in_pool() {
700 threadteam::DispatchKind::InPool
701 } else {
702 threadteam::DispatchKind::Outside
703 }
704}
705
706#[cfg(all(feature = "rayon", feature = "threadteam"))]
708fn host_costs() -> &'static threadteam::DispatchCosts {
709 static HOST: std::sync::OnceLock<threadteam::DispatchCosts> = std::sync::OnceLock::new();
710 HOST.get_or_init(threadteam::DispatchCosts::host_default)
711}
712
713pub(crate) fn par_can_pay_time(par: &crate::Par, estimate: Duration, sections: u32) -> bool {
718 if par.degree() < 2 {
719 return false;
720 }
721 let min = match par {
722 crate::Par::Seq => Duration::MAX,
723 #[cfg(feature = "threadteam")]
724 crate::Par::Team(team) => return par_can_pay_time_team(team, estimate, sections),
725 other => other.shape().min_dispatch_cost(),
726 };
727 pays_at(estimate, sections, min)
728}
729
730#[cfg(feature = "threadteam")]
732pub(crate) fn par_can_pay_time_team(team: &threadteam::TeamExec<'_>, estimate: Duration, sections: u32) -> bool {
733 team.size().get() >= 2
734 && pays_at(estimate, sections, team.dispatch_costs().min_cost(threadteam::DispatchKind::Team))
735}
736
737#[cfg_attr(feature = "par-trace", track_caller)]
740pub(crate) fn par_pays<T: crate::ComplexField>(
741 par: &crate::Par,
742 units: usize,
743 rate: Rate,
744 sections: u32,
745) -> bool {
746 if par.degree() < 2 {
749 return false;
750 }
751 let r = par_pays_time(par, serial_time(units, rate.of::<T>()), sections, par.degree());
752 #[cfg(feature = "par-trace")]
753 trace_decision(par.shape(), par.shape().degree(), units, rate.of::<T>(), sections, r, core::panic::Location::caller());
754 r
755}
756
757#[cfg(feature = "par-trace")]
758fn trace_decision(
759 shape: ParShape,
760 width: usize,
761 units: usize,
762 ns: f64,
763 sections: u32,
764 r: bool,
765 loc: &'static core::panic::Location<'static>,
766) {
767 let d = NonZeroUsize::new(width)
768 .map(|w| shape.dispatch_cost(w).as_secs_f64() * 1e9)
769 .unwrap_or(f64::INFINITY);
770 trace::record(loc, units, sections, units as f64 * ns, d, r);
771}
772
773#[cfg_attr(feature = "par-trace", track_caller)]
778pub(crate) fn join_pays<T: crate::ComplexField>(
779 par: &crate::Par,
780 units: usize,
781 rate: Rate,
782 sections: u32,
783) -> bool {
784 if par.degree() < 2 {
785 return false;
786 }
787 let r = par_pays_time(par, serial_time(units, rate.of::<T>()), sections, 2);
788 #[cfg(feature = "par-trace")]
789 trace_decision(par.shape(), 2, units, rate.of::<T>(), sections, r, core::panic::Location::caller());
790 r
791}
792
793pub(crate) fn halves_pays<T: crate::ComplexField>(
801 par: &crate::Par,
802 n: usize,
803 rate: Rate,
804 power: u32,
805) -> bool {
806 let h = n / 2;
807 join_pays::<T>(par, h.saturating_pow(power).saturating_mul(2), rate, 1)
808}
809
810pub(crate) fn block_apply_estimate<T: crate::ComplexField>(
817 m: usize,
818 n: usize,
819 k: usize,
820) -> Duration {
821 block_apply_estimate_with::<T>(rates(), m, n, k)
822}
823
824pub(crate) fn block_apply_estimate_with<T: crate::ComplexField>(
826 rt: &Rates,
827 m: usize,
828 n: usize,
829 k: usize,
830) -> Duration {
831 let macs = m.saturating_mul(n).saturating_mul(k).saturating_mul(2);
832 let ns = macs as f64 * rt.block_apply_mac.of::<T>() + m.saturating_mul(k) as f64 * rt.householder_rank_one.of::<T>();
834 Duration::from_secs_f64(ns * 1e-9)
835}
836
837pub(crate) fn serial_time(count: usize, ns: f64) -> Duration {
839 Duration::from_secs_f64(count as f64 * ns * 1e-9)
840}
841
842impl ParShape {
843 pub fn dispatch_cost(self, width: NonZeroUsize) -> Duration {
855 let width = width.min(NonZeroUsize::new(self.degree().max(1)).unwrap());
856 if width.get() < 2 {
857 return Duration::MAX;
858 }
859 match self {
860 Self::Seq => Duration::MAX,
861 #[cfg(feature = "rayon")]
862 Self::Rayon(_) => {
863 #[cfg(feature = "threadteam")]
864 {
865 host_costs().cost(rayon_kind(), width)
866 }
867 #[cfg(not(feature = "threadteam"))]
868 {
869 rayon_cost(width)
870 }
871 },
872 #[cfg(feature = "threadteam")]
873 Self::Team(_, costs) => costs.cost(threadteam::DispatchKind::Team, width),
874 }
875 }
876
877 pub fn min_dispatch_cost(self) -> Duration {
881 if self.degree() < 2 {
882 return Duration::MAX;
883 }
884 match self {
885 Self::Seq => Duration::MAX,
886 #[cfg(feature = "rayon")]
887 Self::Rayon(_) => {
888 #[cfg(feature = "threadteam")]
889 {
890 host_costs().min_cost(rayon_kind())
891 }
892 #[cfg(not(feature = "threadteam"))]
893 {
894 min_rayon_cost()
895 }
896 },
897 #[cfg(feature = "threadteam")]
898 Self::Team(_, costs) => costs.min_cost(threadteam::DispatchKind::Team),
899 }
900 }
901
902 pub fn pays(self, estimate: Duration, sections: u32, width: NonZeroUsize) -> bool {
906 pays_at(estimate, sections, self.dispatch_cost(width))
907 }
908
909 pub fn pays_full(self, estimate: Duration, sections: u32) -> bool {
911 match NonZeroUsize::new(self.degree()) {
912 Some(width) => self.pays(estimate, sections, width),
913 None => false,
914 }
915 }
916}
917
918#[cfg(test)]
919mod tests {
920 use super::*;
921 use crate::c64;
922
923 fn nz(n: usize) -> NonZeroUsize {
924 NonZeroUsize::new(n).unwrap()
925 }
926
927 #[test]
928 fn seq_never_pays() {
929 assert_eq!(ParShape::Seq.dispatch_cost(nz(4)), Duration::MAX);
930 assert!(!ParShape::Seq.pays(Duration::from_secs(10), 1, nz(4)));
931 assert!(!ParShape::Seq.pays_full(Duration::from_secs(10), 1));
932 }
933
934 #[test]
935 #[cfg(feature = "threadteam")]
936 fn min_dispatch_cost_is_the_cheapest_width() {
937 use threadteam::{DispatchCosts, DispatchKind, WidthCosts};
938 let us = Duration::from_micros;
939 for t in [[us(9), us(5), us(3)], [us(4), us(4), us(4)], [us(2), us(7), us(12)]] {
941 let costs = DispatchCosts::for_class(threadteam::HostClass::Generic).with(DispatchKind::Team, WidthCosts::new(t));
942 let s = ParShape::Team(nz(6), costs);
943 let want = (2..=6).map(|w| s.dispatch_cost(nz(w))).min().unwrap();
944 assert_eq!(s.min_dispatch_cost(), want, "{t:?}");
945 }
946 assert_eq!(ParShape::Seq.min_dispatch_cost(), Duration::MAX);
947 assert_eq!(ParShape::Team(nz(1), DispatchCosts::for_class(threadteam::HostClass::Generic)).min_dispatch_cost(), Duration::MAX);
948 }
949
950 #[test]
951 fn pays_is_the_multiple_of_the_cost() {
952 let d = Duration::from_micros(2);
953 assert!(pays_at(Duration::from_micros(6), 1, d));
954 assert!(!pays_at(Duration::from_nanos(5999), 1, d));
955 assert!(!pays_at(Duration::from_micros(11), 2, d));
956 assert!(pays_at(Duration::from_micros(12), 2, d));
957 assert!(!pays_at(Duration::from_secs(1), 1, Duration::MAX));
958 }
959
960 #[cfg(feature = "rayon")]
961 #[test]
962 fn rayon_uses_the_outside_cost() {
963 let s = ParShape::Rayon(nz(6));
964 assert_eq!(s.dispatch_cost(nz(1)), Duration::MAX);
965 let d6 = s.dispatch_cost(nz(6));
966 assert!(d6 > Duration::ZERO && d6 < Duration::MAX);
969 assert_eq!(s.dispatch_cost(nz(60)), d6);
971 assert!(!s.pays(d6, 1, nz(6)));
973 assert!(s.pays(d6 * 4, 1, nz(6)));
974 }
975
976 #[cfg(feature = "threadteam")]
977 #[test]
978 fn team_shape_carries_the_costs_of_its_token() {
979 use threadteam::{DispatchCosts, DispatchKind, WidthCosts};
980 let costs = DispatchCosts::for_class(threadteam::HostClass::Generic)
981 .with(DispatchKind::Team, WidthCosts::flat(Duration::from_micros(5)));
982 let s = ParShape::Team(nz(4), costs);
983 assert_eq!(s.dispatch_cost(nz(4)), Duration::from_micros(5));
984 assert_eq!(s.dispatch_cost(nz(1)), Duration::MAX);
985 assert!(s.pays(Duration::from_micros(15), 1, nz(4)));
986 assert!(!s.pays(Duration::from_micros(14), 1, nz(4)));
987 }
988
989 #[test]
990 fn gemm_rates_are_per_dtype() {
991 assert_eq!(gemm_ns_per_mac::<f64>(), ns_per_mac::f64());
992 assert_eq!(gemm_ns_per_mac::<crate::c64>(), ns_per_mac::c64());
993 assert_eq!(gemm_ns_per_mac::<f32>(), ns_per_mac::f32());
994 assert_eq!(gemm_ns_per_mac::<crate::c32>(), ns_per_mac::c32());
995 assert!(ns_per_mac::c64() > 4.0 * ns_per_mac::f64());
997 }
998
999 fn floor_units<T: crate::ComplexField>(
1001 shape: ParShape,
1002 r: Rate,
1003 sections: u32,
1004 ) -> usize {
1005 let pays = |u: usize| {
1006 ParThreshold::Derived.region_pays::<T>(shape, u, r, sections)
1007 };
1008 let mut hi = 1usize;
1009 while !pays(hi) {
1010 hi *= 2;
1011 assert!(hi < 1 << 40, "never pays");
1012 }
1013 let mut lo = hi / 2;
1014 while hi - lo > 1 {
1015 let mid = (lo + hi) / 2;
1016 if pays(mid) {
1017 hi = mid;
1018 } else {
1019 lo = mid;
1020 }
1021 }
1022 hi
1023 }
1024
1025 fn gates(rt: &Rates) -> [(&'static str, Rate, u32); 10] {
1027 [
1028 ("lu_full_pivot", rt.lu_full_pivot, 1),
1029 ("bunch_kaufman", rt.bunch_kaufman, 1),
1030 ("qr_col_pivot", rt.qr_col_pivot, 1),
1031 ("bidiag", rt.bidiag, 2),
1032 ("hessenberg", rt.hessenberg, 2),
1033 ("tridiag", rt.tridiag, 1),
1034 ("matvec", rt.matvec, 1),
1035 ("householder_rank_one", rt.householder_rank_one, 1),
1036 ("gemm_mac", rt.gemm_mac, 1),
1037 ("sparse_matmul", rt.sparse_matmul, 1),
1038 ]
1039 }
1040
1041 const M5: &Rates = &Rates::AARCH64_APPLE;
1044
1045 fn entries(rt: &Rates) -> [(&'static str, Rate); 16] {
1047 [
1048 ("gemm_mac", rt.gemm_mac),
1049 ("gemm_single", rt.gemm_single),
1050 ("block_apply_mac", rt.block_apply_mac),
1051 ("lu_full_pivot", rt.lu_full_pivot),
1052 ("bunch_kaufman", rt.bunch_kaufman),
1053 ("qr_col_pivot", rt.qr_col_pivot),
1054 ("bidiag", rt.bidiag),
1055 ("hessenberg", rt.hessenberg),
1056 ("tridiag", rt.tridiag),
1057 ("matvec", rt.matvec),
1058 ("bidiag_svd_dc", rt.bidiag_svd_dc),
1059 ("tridiag_evd_dc", rt.tridiag_evd_dc),
1060 ("triangular_inverse", rt.triangular_inverse),
1061 ("sparse_matmul", rt.sparse_matmul),
1062 ("householder_rank_one", rt.householder_rank_one),
1063 ("lu_swap", rt.lu_swap),
1064 ]
1065 }
1066
1067 #[test]
1068 fn every_class_resolves_to_its_table() {
1069 assert_eq!(Rates::for_class(HostClass::X86_64Avx512), &Rates::X86_64_AVX512);
1070 assert_eq!(Rates::for_class(HostClass::Aarch64Apple), &Rates::AARCH64_APPLE);
1071 for c in [HostClass::X86_64Avx2, HostClass::X86_64Scalar, HostClass::Aarch64, HostClass::Generic] {
1072 assert_eq!(Rates::for_class(c), &Rates::GENERIC, "{}", c.name());
1073 }
1074 assert_eq!(Rates::X86_64_AVX512.gemm_mac, Rate { real: 0.0641, cplx: 0.2565 });
1076 assert_eq!(Rates::X86_64_AVX512.gemm_single, Rate { real: 0.0320, cplx: 0.1283 });
1077 assert_eq!(Rates::AARCH64_APPLE.gemm_mac, Rate { real: 0.037, cplx: 0.165 });
1078 assert_eq!(Rates::AARCH64_APPLE.bidiag_svd_dc, Rate { real: 53.0, cplx: 53.0 });
1079 let g = Rates::X86_64_AVX512;
1081 assert!((g.gemm_single.real / (g.gemm_mac.real / 2.0) - 1.0).abs() < 0.01);
1082 assert!((g.gemm_single.cplx / (g.gemm_mac.cplx / 2.0) - 1.0).abs() < 0.01);
1083 }
1084
1085 #[test]
1086 fn generic_is_the_per_entry_minimum_of_the_measured_rows() {
1087 let (a, b, g) = (entries(&Rates::X86_64_AVX512), entries(&Rates::AARCH64_APPLE), entries(&Rates::GENERIC));
1088 for i in 0..a.len() {
1089 let (name, ga) = (a[i].0, a[i].1);
1090 assert_eq!(g[i].1.real, ga.real.min(b[i].1.real), "{name} real");
1091 assert_eq!(g[i].1.cplx, ga.cplx.min(b[i].1.cplx), "{name} cplx");
1092 assert!(g[i].1.real <= ga.real && g[i].1.real <= b[i].1.real, "{name}");
1093 }
1094 }
1095
1096 #[test]
1097 fn rows_are_positive() {
1098 for c in HostClass::ALL {
1099 for (name, r) in entries(Rates::for_class(c)) {
1100 assert!(r.real > 0.0 && r.cplx > 0.0, "{} {name}", c.name());
1101 }
1102 }
1103 }
1104
1105 #[test]
1106 fn detection_is_stable_and_the_host_row_is_its_class() {
1107 let c = HostClass::detect();
1108 for _ in 0..4 {
1109 assert_eq!(HostClass::detect(), c);
1110 }
1111 assert_eq!(HostClass::detect_uncached(), c);
1112 assert_eq!(rates(), Rates::for_class(c));
1113 assert_eq!(rate::bidiag(), rates().bidiag);
1114 #[cfg(all(target_arch = "x86_64", feature = "std"))]
1115 if std::arch::is_x86_feature_detected!("avx512f") {
1116 assert_eq!(c, HostClass::X86_64Avx512);
1117 }
1118 #[cfg(all(target_arch = "aarch64", target_os = "macos"))]
1119 assert_eq!(c, HostClass::Aarch64Apple);
1120 }
1121
1122 #[test]
1123 fn outside_rows_follow_the_classes() {
1124 let us = |c, w: usize| outside_cost_of(c, nz(w)).as_nanos();
1125 assert_eq!((us(HostClass::X86_64Avx512, 2), us(HostClass::X86_64Avx512, 4), us(HostClass::X86_64Avx512, 6)), (9_700, 10_100, 8_700));
1126 assert_eq!((us(HostClass::Aarch64Apple, 2), us(HostClass::Aarch64Apple, 4), us(HostClass::Aarch64Apple, 6)), (14_000, 18_000, 19_000));
1127 assert_eq!(us(HostClass::Aarch64Apple, 3), 16_000);
1128 assert_eq!(us(HostClass::Aarch64Apple, 5), 18_500);
1129 for c in [HostClass::X86_64Avx2, HostClass::X86_64Scalar, HostClass::Aarch64, HostClass::Generic] {
1130 for w in [1, 2, 3, 6, 12] {
1131 assert_eq!(us(c, w), 40_000);
1132 }
1133 }
1134 assert_eq!(us(HostClass::Aarch64Apple, 12), 19_000);
1136 }
1137
1138 #[test]
1139 fn min_rayon_cost_is_the_cheapest_width() {
1140 for c in HostClass::ALL {
1141 let widths = |f: fn(HostClass, NonZeroUsize) -> Duration| (2..=12).map(|w| f(c, nz(w))).min().unwrap();
1142 assert_eq!(widths(outside_cost_of), Duration::from_nanos(outside_anchors(c).into_iter().min().unwrap()));
1143 assert_eq!(widths(in_pool_cost_of), Duration::from_nanos(in_pool_anchors(c).into_iter().min().unwrap()));
1144 }
1145 }
1146
1147 #[test]
1148 fn in_pool_rows_follow_the_classes() {
1149 let us = |c, w: usize| in_pool_cost_of(c, nz(w)).as_nanos();
1150 assert_eq!((us(HostClass::X86_64Avx512, 2), us(HostClass::X86_64Avx512, 4), us(HostClass::X86_64Avx512, 6)), (3_000, 3_600, 3_500));
1151 assert_eq!((us(HostClass::Aarch64Apple, 2), us(HostClass::Aarch64Apple, 4), us(HostClass::Aarch64Apple, 6)), (6_000, 11_000, 5_500));
1152 assert_eq!(us(HostClass::Generic, 8), 15_000);
1153 for c in HostClass::ALL {
1155 for w in 2..=12 {
1156 assert!(in_pool_cost_of(c, nz(w)) <= outside_cost_of(c, nz(w)), "{} width {w}", c.name());
1157 }
1158 }
1159 }
1160
1161 #[cfg(feature = "rayon")]
1165 #[test]
1166 fn a_held_spindle_group_is_a_worker_of_the_pool() {
1167 let s = ParShape::Rayon(nz(8));
1168 let outside = s.dispatch_cost(nz(8));
1169 assert!(!caller_in_pool());
1170 let held = spindle::with_lock(8, || (caller_in_pool(), s.dispatch_cost(nz(8)), s.min_dispatch_cost()));
1171 assert!(held.0);
1172 assert!(held.1 <= outside && held.2 <= s.min_dispatch_cost());
1173 if HostClass::detect() == HostClass::X86_64Avx512 {
1174 assert!(held.1 < outside);
1175 }
1176 }
1177
1178 #[cfg(feature = "threadteam")]
1181 #[test]
1182 fn host_class_and_outside_rows_agree_with_threadteam() {
1183 use threadteam::{DispatchCosts, DispatchKind};
1184 let tt = |c: HostClass| match c {
1185 HostClass::X86_64Avx512 => threadteam::HostClass::X86_64Avx512,
1186 HostClass::X86_64Avx2 => threadteam::HostClass::X86_64Avx2,
1187 HostClass::X86_64Scalar => threadteam::HostClass::X86_64Scalar,
1188 HostClass::Aarch64Apple => threadteam::HostClass::Aarch64Apple,
1189 HostClass::Aarch64 => threadteam::HostClass::Aarch64,
1190 HostClass::Generic => threadteam::HostClass::Generic,
1191 };
1192 assert_eq!(tt(HostClass::detect()), threadteam::HostClass::detect());
1193 for c in HostClass::ALL {
1194 assert_eq!(c.name(), tt(c).name());
1195 let costs = DispatchCosts::for_class(tt(c));
1196 for w in 2..=12 {
1197 assert_eq!(outside_cost_of(c, nz(w)), costs.cost(DispatchKind::Outside, nz(w)), "{} width {w}", c.name());
1198 assert_eq!(in_pool_cost_of(c, nz(w)), costs.cost(DispatchKind::InPool, nz(w)), "{} width {w}", c.name());
1199 }
1200 }
1201 }
1202
1203 #[test]
1204 fn fixed_and_never_keep_their_meaning() {
1205 let s = ParShape::Rayon(nz(6));
1206 let r = rate::lu_full_pivot();
1207 assert!(!ParThreshold::Never.region_pays::<f64>(s, usize::MAX, r, 1));
1208 assert!(ParThreshold::Fixed(100).region_pays::<f64>(ParShape::Seq, 100, r, 1));
1209 assert!(!ParThreshold::Fixed(100).region_pays::<f64>(s, 99, r, 1));
1210 assert!(ParThreshold::Fixed(100).region_pays::<f64>(s, 100, r, 1));
1211 assert_eq!(ParThreshold::from(0), ParThreshold::Derived);
1213 assert_eq!(ParThreshold::from(usize::MAX), ParThreshold::Never);
1214 assert_eq!(ParThreshold::from(7), ParThreshold::Fixed(7));
1215 assert_eq!(usize::from(ParThreshold::Fixed(0)), 1);
1216 }
1217
1218 #[test]
1219 fn derived_gates_never_pay_on_a_sequential_shape() {
1220 for (_, r, sections) in gates(rates()) {
1221 assert!(!ParThreshold::Derived.region_pays::<f64>(
1222 ParShape::Seq,
1223 usize::MAX / 2,
1224 r,
1225 sections
1226 ));
1227 assert!(!ParThreshold::Derived.region_pays::<c64>(
1228 ParShape::Seq,
1229 usize::MAX / 2,
1230 r,
1231 sections
1232 ));
1233 }
1234 }
1235
1236 #[cfg(feature = "threadteam")]
1237 #[test]
1238 fn derived_floors_on_a_team_follow_the_rule() {
1239 use threadteam::DispatchCosts;
1240 let team = ParShape::Team(nz(6), DispatchCosts::for_class(threadteam::HostClass::Aarch64Apple));
1241 let d = team.dispatch_cost(nz(6)).as_secs_f64() * 1e9;
1242 for (name, r, sections) in gates(M5) {
1243 for (ty, ns) in [("f64", r.real), ("c64", r.cplx)] {
1244 let floor = if ty == "f64" {
1245 floor_units::<f64>(team, r, sections)
1246 } else {
1247 floor_units::<c64>(team, r, sections)
1248 };
1249 let want = PAYS * f64::from(sections) * d / ns;
1251 assert!(
1252 (floor as f64 - want).abs() <= 4.0 + want * 2e-3,
1253 "{name} {ty}: floor {floor}, rule {want}"
1254 );
1255 }
1256 }
1257 let pays = |u: usize, r: Rate, f: u32| {
1259 ParThreshold::Derived.region_pays::<f64>(team, u, r, f)
1260 };
1261 assert!(!pays(100 * 100, M5.lu_full_pivot, 1));
1263 assert!(pays(200 * 200, M5.lu_full_pivot, 1));
1264 assert!(floor_units::<f64>(team, M5.bidiag, 2) >= 2 * floor_units::<f64>(team, M5.bidiag, 1) - 2);
1266 assert!(pays(256 * 256, M5.matvec, 1));
1268 assert!(!pays(64 * 64, M5.matvec, 1));
1269 let lu = floor_units::<f64>(team, M5.lu_full_pivot, 1);
1272 let bidiag = floor_units::<f64>(team, M5.bidiag, 1);
1273 assert!(lu > bidiag);
1274 }
1275
1276 #[cfg(all(feature = "rayon", target_os = "macos", target_arch = "aarch64"))]
1277 #[test]
1278 fn derived_floors_outside_the_pool_are_larger_than_on_a_team() {
1279 let rayon = ParShape::Rayon(nz(6));
1280 for (name, r, sections) in gates(M5) {
1281 let outside = floor_units::<f64>(rayon, r, sections);
1282 #[cfg(feature = "threadteam")]
1283 {
1284 use threadteam::DispatchCosts;
1285 let team = ParShape::Team(nz(6), DispatchCosts::for_class(threadteam::HostClass::Aarch64Apple));
1286 let t = floor_units::<f64>(team, r, sections);
1287 assert!(outside > 8 * t, "{name}: outside {outside}, team {t}");
1288 }
1289 let want = PAYS * f64::from(sections) * 19_000.0 / r.real;
1291 assert!((outside as f64 - want).abs() <= 4.0 + want * 2e-3, "{name}");
1292 }
1293 assert!(floor_units::<f64>(rayon, M5.bidiag, 2) > 192 * 256);
1296 }
1297
1298 #[cfg(feature = "rayon")]
1299 #[test]
1300 fn par_pays_reads_the_degree_of_the_token() {
1301 let seq = crate::Par::Seq;
1302 assert!(!par_pays::<f64>(&seq, usize::MAX / 2, rate::matvec(), 1));
1303 let par = crate::Par::rayon(6);
1304 assert!(par_pays::<f64>(&par, 10_000_000, rate::matvec(), 1));
1305 assert!(!par_pays::<f64>(&par, 1_000, rate::matvec(), 1));
1306 let one = crate::Par::rayon(1);
1307 assert!(!par_pays::<f64>(&one, 10_000_000, rate::matvec(), 1));
1308 }
1309
1310 #[test]
1311 fn rates_are_per_dtype() {
1312 let r = rate::bidiag();
1313 assert_eq!(r.of::<f64>(), r.real);
1314 assert_eq!(r.of::<c64>(), r.cplx);
1315 assert_eq!(r.of::<f32>(), r.real / 2.0);
1316 assert_eq!(r.of::<crate::c32>(), r.cplx / 2.0);
1317 assert_eq!(rate::gemm_mac().of::<f64>(), ns_per_mac::f64());
1318 assert_eq!(rate::gemm_mac().of::<c64>(), ns_per_mac::c64());
1319 }
1320
1321 #[cfg(feature = "rayon")]
1325 mod forced_parallel {
1326 use super::*;
1327 use crate::dyn_stack::{MemBuffer, MemStack};
1328 use crate::internal_prelude::*;
1329 use crate::linalg::cholesky::lblt::factor as bk;
1330 use crate::linalg::evd::{hessenberg, tridiag};
1331 use crate::linalg::lu::full_pivoting::factor as fplu;
1332 use crate::linalg::qr::col_pivoting::factor as cpqr;
1333 use crate::linalg::svd::bidiag;
1334 use crate::{Auto, Mat, MatRef, Par};
1335
1336 const N: usize = 96;
1337
1338 fn fill<T: crate::ComplexField>(n: usize, seed: usize) -> Mat<T> {
1339 Mat::from_fn(n, n, |i, j| {
1340 let h = (i * 2654435761usize)
1341 .wrapping_add(j * 40503)
1342 .wrapping_add(seed * 977)
1343 % 1009;
1344 from_f64::<T>(h as f64 / 504.0 - 1.0)
1345 })
1346 }
1347 fn close<T: crate::ComplexField<Real = f64>>(a: MatRef<'_, T>, b: MatRef<'_, T>) {
1348 let mut err = 0.0f64;
1349 for j in 0..a.ncols() {
1350 for i in 0..a.nrows() {
1351 err = err.max((a[(i, j)].clone() - b[(i, j)].clone()).abs());
1352 }
1353 }
1354 core::assert!(err <= 1e-9, "error {err}");
1355 }
1356
1357 fn check<T: crate::ComplexField<Real = f64>>() {
1358 let a0 = fill::<T>(N, 1);
1359 let sym = &a0 + a0.adjoint();
1360 let lu = |par: Par| {
1362 let mut a = a0.clone();
1363 let mut p = vec![0usize; 4 * N];
1364 let (r, rest) = p.split_at_mut(N);
1365 let (ri, rest) = rest.split_at_mut(N);
1366 let (c, ci) = rest.split_at_mut(N);
1367 let mut params: fplu::FullPivLuParams = Auto::<T>::auto();
1368 params.par_threshold = ParThreshold::Fixed(1);
1369 let mut mem = MemBuffer::new(fplu::lu_in_place_scratch::<usize, T>(
1370 N,
1371 N,
1372 par.shape(),
1373 params.into(),
1374 ));
1375 fplu::lu_in_place(
1376 a.as_mut(), r, ri, c, ci, par, MemStack::new(&mut mem), params.into(),
1377 );
1378 a
1379 };
1380 close(lu(Par::Seq).as_ref(), lu(Par::rayon(4)).as_ref());
1381 let bunch = |par: Par| {
1382 let mut a = sym.clone();
1383 let mut perm = vec![0usize; N];
1384 let mut perm_inv = vec![0usize; N];
1385 let mut sub = Mat::<T>::zeros(N, 1);
1386 let mut params: bk::LbltParams = Auto::<T>::auto();
1387 params.pivoting = bk::PivotingStrategy::Full;
1388 params.par_threshold = ParThreshold::Fixed(1);
1389 let mut mem = MemBuffer::new(bk::cholesky_in_place_scratch::<usize, T>(
1390 N,
1391 par.shape(),
1392 params.into(),
1393 ));
1394 bk::cholesky_in_place(
1395 a.as_mut(),
1396 sub.as_mut().col_mut(0).as_diagonal_mut(),
1397 &mut perm,
1398 &mut perm_inv,
1399 par,
1400 MemStack::new(&mut mem),
1401 params.into(),
1402 );
1403 a
1404 };
1405 close(bunch(Par::Seq).as_ref(), bunch(Par::rayon(4)).as_ref());
1406 let qr = |par: Par| {
1407 let mut a = a0.clone();
1408 let mut h = Mat::<T>::zeros(16, N);
1409 let mut p = vec![0usize; N];
1410 let mut pi = vec![0usize; N];
1411 let mut params: cpqr::ColPivQrParams = Auto::<T>::auto();
1412 params.par_threshold = ParThreshold::Fixed(1);
1413 let mut mem = MemBuffer::new(cpqr::qr_in_place_scratch::<usize, T>(
1414 N,
1415 N,
1416 16,
1417 par.shape(),
1418 params.into(),
1419 ));
1420 cpqr::qr_in_place(
1421 a.as_mut(), h.as_mut(), &mut p, &mut pi, par, MemStack::new(&mut mem), params.into(),
1422 );
1423 a
1424 };
1425 close(qr(Par::Seq).as_ref(), qr(Par::rayon(4)).as_ref());
1426 let bid = |par: Par| {
1427 let mut a = a0.clone();
1428 let mut hl = Mat::<T>::zeros(8, N);
1429 let mut hr = Mat::<T>::zeros(8, N - 1);
1430 let mut params: bidiag::BidiagParams = Auto::<T>::auto();
1431 params.par_threshold = ParThreshold::Fixed(1);
1432 let mut mem = MemBuffer::new(bidiag::bidiag_in_place_scratch::<T>(
1433 N,
1434 N,
1435 par.shape(),
1436 params.into(),
1437 ));
1438 bidiag::bidiag_in_place(
1439 a.as_mut(), hl.as_mut(), hr.as_mut(), par, MemStack::new(&mut mem), params.into(),
1440 );
1441 a
1442 };
1443 close(bid(Par::Seq).as_ref(), bid(Par::rayon(4)).as_ref());
1444 let hess = |par: Par| {
1445 let mut a = a0.clone();
1446 let mut h = Mat::<T>::zeros(8, N - 1);
1447 let mut params: hessenberg::HessenbergParams = Auto::<T>::auto();
1448 params.par_threshold = ParThreshold::Fixed(1);
1449 params.blocking_threshold = usize::MAX;
1450 let mut mem = MemBuffer::new(hessenberg::hessenberg_in_place_scratch::<T>(
1451 N,
1452 8,
1453 par.shape(),
1454 params.into(),
1455 ));
1456 hessenberg::hessenberg_in_place(
1457 a.as_mut(), h.as_mut(), par, MemStack::new(&mut mem), params.into(),
1458 );
1459 a
1460 };
1461 close(hess(Par::Seq).as_ref(), hess(Par::rayon(4)).as_ref());
1462 let tri = |par: Par| {
1463 let mut a = sym.clone();
1464 let mut h = Mat::<T>::zeros(8, N - 1);
1465 let mut params: tridiag::TridiagParams = Auto::<T>::auto();
1466 params.par_threshold = ParThreshold::Fixed(1);
1467 let mut mem = MemBuffer::new(tridiag::tridiag_in_place_scratch::<T>(
1468 N,
1469 par.shape(),
1470 params.into(),
1471 ));
1472 tridiag::tridiag_in_place(
1473 a.as_mut(), h.as_mut(), par, MemStack::new(&mut mem), params.into(),
1474 );
1475 a
1476 };
1477 close(tri(Par::Seq).as_ref(), tri(Par::rayon(4)).as_ref());
1478 }
1479
1480 #[test]
1481 fn forced_parallel_steps_equal_serial_f64() {
1482 check::<f64>();
1483 }
1484
1485 #[test]
1486 fn forced_parallel_steps_equal_serial_c64() {
1487 check::<c64>();
1488 }
1489 }
1490
1491 #[test]
1493 fn halves_join_gates_follow_the_rule() {
1494 let seq = crate::Par::Seq;
1495 for r in [rate::bidiag_svd_dc(), rate::tridiag_evd_dc(), rate::triangular_inverse()] {
1496 assert!(!halves_pays::<f64>(&seq, 1 << 20, r, 2));
1497 }
1498 #[cfg(feature = "rayon")]
1499 {
1500 let par = crate::Par::rayon(6);
1501 for r in [rate::bidiag_svd_dc(), rate::tridiag_evd_dc(), rate::triangular_inverse()] {
1502 let power = if r == rate::triangular_inverse() { 3 } else { 2 };
1504 let first = (2..4096).find(|&n| halves_pays::<f64>(&par, n, r, power)).expect("never pays");
1505 assert!(first > 8, "{first}");
1506 assert!(!halves_pays::<f64>(&par, first - 1, r, power));
1507 assert!(halves_pays::<f64>(&par, first, r, power));
1508 assert!(halves_pays::<f64>(&par, 2 * first, r, power));
1509 let h = first / 2;
1511 let work_ns = 2.0 * (h as f64).powi(power as i32) * r.real;
1512 let want = PAYS * par.shape().dispatch_cost(nz(2)).as_secs_f64() * 1e9;
1513 assert!(work_ns >= want * 0.999, "{work_ns} {want}");
1514 }
1515 }
1516 }
1517
1518 #[cfg(feature = "threadteam")]
1525 #[test]
1526 fn block_apply_split_starts_where_the_estimate_pays() {
1527 use threadteam::DispatchCosts;
1528 let team = ParShape::Team(nz(6), DispatchCosts::for_class(threadteam::HostClass::Aarch64Apple));
1529 let pays = |m: usize| {
1530 let (n, k) = (8, m - 8);
1531 pays_at(block_apply_estimate_with::<f64>(M5, m, n, k), 1, team.dispatch_cost(nz(6)))
1532 };
1533 assert!(!pays(72));
1534 assert!(pays(80));
1535 let gemm_only = |m: usize| {
1537 pays_at(serial_time(2 * m * 8 * (m - 8), M5.gemm_mac.real), 1, team.dispatch_cost(nz(6)))
1538 };
1539 assert!(!gemm_only(88) && gemm_only(96));
1540 let c = block_apply_estimate_with::<c64>(M5, 100, 8, 92).as_secs_f64() * 1e9;
1542 let want = 2.0 * 100.0 * 8.0 * 92.0 * M5.block_apply_mac.cplx + 100.0 * 92.0 * M5.householder_rank_one.cplx;
1543 assert!((c - want).abs() < 2.0, "{c} {want}");
1544 }
1545
1546 #[cfg(feature = "threadteam")]
1547 #[test]
1548 fn join_gates_on_a_team_follow_the_width_two_cost() {
1549 use threadteam::DispatchCosts;
1550 let team = ParShape::Team(nz(6), DispatchCosts::for_class(threadteam::HostClass::Aarch64Apple));
1551 let rayon = ParShape::Rayon(nz(6));
1552 for (name, r, sections) in gates(M5) {
1553 let t = floor_units::<f64>(team, r, sections);
1554 let o = floor_units::<f64>(rayon, r, sections);
1555 assert!(t < o, "{name}: team {t}, outside {o}");
1556 }
1557 let d2 = team.dispatch_cost(nz(2));
1559 assert_eq!(d2, Duration::from_micros(3));
1560 let par = |ns: f64| pays_at(Duration::from_secs_f64(ns * 1e-9), 1, d2);
1561 assert!(par(9000.0) && !par(8900.0));
1562 for (r, power, yes, no) in [
1566 (M5.bidiag_svd_dc, 2u32, 128usize, 8usize),
1567 (M5.tridiag_evd_dc, 2, 64, 8),
1568 (M5.triangular_inverse, 3, 256, 64),
1569 ] {
1570 let pays = |n: usize| {
1571 let h = n / 2;
1572 pays_at(
1573 serial_time(2 * h.pow(power), r.real),
1574 1,
1575 d2,
1576 )
1577 };
1578 assert!(pays(yes));
1579 assert!(!pays(no));
1580 }
1581 assert!(ParThreshold::Derived.region_pays::<f64>(team, 1_000_000, M5.sparse_matmul, 1));
1583 assert!(!ParThreshold::Derived.region_pays::<f64>(team, 1_000, M5.sparse_matmul, 1));
1584 }
1585}
1586
1587#[cfg(feature = "par-trace")]
1588#[doc(hidden)]
1589pub fn trace_pub(
1590 units: usize,
1591 sections: u32,
1592 estimate: Duration,
1593 dispatch: Duration,
1594 r: bool,
1595 loc: &'static core::panic::Location<'static>,
1596) {
1597 trace::record(loc, units, sections, estimate.as_secs_f64() * 1e9, dispatch.as_secs_f64() * 1e9, r);
1598}
1599
1600#[cfg(all(test, feature = "rayon", feature = "sparse-linalg"))]
1601mod join_results {
1602 use crate::dyn_stack::{MemBuffer, MemStack};
1605 use crate::linalg::evd::{self, ComputeEigenvectors};
1606 use crate::linalg::svd;
1607 use crate::linalg::triangular_inverse::{invert_lower_triangular, invert_unit_lower_triangular};
1608 use crate::{Accum, Mat, Par, Spec};
1609
1610 struct Rng(u64);
1611 impl Rng {
1612 fn next(&mut self) -> f64 {
1613 self.0 ^= self.0 << 13;
1614 self.0 ^= self.0 >> 7;
1615 self.0 ^= self.0 << 17;
1616 (self.0 >> 11) as f64 / (1u64 << 53) as f64 - 0.5
1617 }
1618 }
1619
1620 fn close(a: &[f64], b: &[f64], tol: f64) {
1621 let scale = a.iter().fold(1.0f64, |m, x| m.max(x.abs()));
1622 for (x, y) in a.iter().zip(b) {
1623 assert!((x - y).abs() <= tol * scale, "{x} {y}");
1624 }
1625 }
1626
1627 #[test]
1628 fn triangular_inverse_agrees_with_serial() {
1629 let n = 512;
1630 let mut r = Rng(3);
1631 let mut src = Mat::<f64>::from_fn(n, n, |_, _| r.next() * 0.1);
1632 for i in 0..n {
1633 src[(i, i)] = 2.0 + src[(i, i)].abs();
1634 }
1635 for unit in [false, true] {
1636 let mut a = Mat::<f64>::zeros(n, n);
1637 let mut b = Mat::<f64>::zeros(n, n);
1638 for (dst, par) in [(&mut a, Par::Seq), (&mut b, Par::rayon(4))] {
1639 if unit {
1640 invert_unit_lower_triangular(dst.as_mut(), src.as_ref(), par);
1641 } else {
1642 invert_lower_triangular(dst.as_mut(), src.as_ref(), par);
1643 }
1644 }
1645 for j in 0..n {
1646 close(
1647 &(0..n).map(|i| a[(i, j)]).collect::<Vec<_>>(),
1648 &(0..n).map(|i| b[(i, j)]).collect::<Vec<_>>(),
1649 1e-9,
1650 );
1651 }
1652 }
1653 }
1654
1655 #[test]
1656 fn bidiagonal_svd_agrees_with_serial() {
1657 let n = 512;
1658 let mut r = Rng(5);
1659 let d0: Vec<f64> = (0..n).map(|_| r.next() + 1.0).collect();
1660 let e: Vec<f64> = (0..n - 1).map(|_| r.next()).collect();
1661 let mut out = vec![];
1662 for par in [Par::Seq, Par::rayon(4)] {
1663 let mut d = d0.clone();
1664 let mut u = Mat::<f64>::zeros(n, n);
1665 let mut v = Mat::<f64>::zeros(n, n);
1666 let mut mem = MemBuffer::new(svd::bidiagonal_svd_scratch::<f64>(
1667 n,
1668 par.shape(),
1669 Spec::default(),
1670 ));
1671 svd::bidiagonal_svd(
1672 &mut d,
1673 &e,
1674 true,
1675 u.as_mut(),
1676 v.as_mut(),
1677 par,
1678 MemStack::new(&mut mem),
1679 Spec::default(),
1680 )
1681 .unwrap();
1682 out.push(d);
1683 }
1684 close(&out[0], &out[1], 1e-10);
1685 }
1686
1687 #[test]
1688 fn tridiagonal_evd_agrees_with_serial() {
1689 let n = 512;
1690 let mut r = Rng(6);
1691 let d0: Vec<f64> = (0..n).map(|_| r.next()).collect();
1692 let mut e0: Vec<f64> = (0..n - 1).map(|_| r.next()).collect();
1693 e0.push(0.0);
1694 let mut out = vec![];
1695 for par in [Par::Seq, Par::rayon(4)] {
1696 let mut u = Mat::<f64>::zeros(n, n);
1697 let mut s = Mat::<f64>::zeros(n, 1);
1698 let mut mem = MemBuffer::new(evd::tridiagonal_self_adjoint_evd_scratch::<f64>(
1699 n,
1700 ComputeEigenvectors::Yes,
1701 par.shape(),
1702 Spec::default(),
1703 ));
1704 let d = crate::col::ColRef::from_slice(&d0);
1705 let e = crate::col::ColRef::from_slice(&e0);
1706 evd::tridiagonal_self_adjoint_evd(
1707 d.as_diagonal(),
1708 e.as_diagonal(),
1709 s.as_mut().col_mut(0).as_diagonal_mut(),
1710 Some(u.as_mut()),
1711 par,
1712 MemStack::new(&mut mem),
1713 Spec::default(),
1714 )
1715 .unwrap();
1716 out.push((0..n).map(|i| s[(i, 0)]).collect::<Vec<_>>());
1717 }
1718 close(&out[0], &out[1], 1e-10);
1719 }
1720
1721 #[test]
1722 fn sparse_matmul_agrees_with_serial() {
1723 use crate::sparse::SparseColMat;
1724 use crate::sparse::Triplet;
1725 use crate::sparse::linalg::matmul::sparse_sparse_matmul;
1726 let n = 6000;
1727 let mk = |seed: u64| {
1728 let mut r = Rng(seed);
1729 let mut t = vec![];
1730 for j in 0..n {
1731 for _ in 0..24 {
1732 let i = ((r.next() + 0.5) * n as f64) as usize % n;
1733 t.push(Triplet::new(i, j, r.next()));
1734 }
1735 }
1736 SparseColMat::<usize, f64>::try_new_from_triplets(n, n, &t).unwrap()
1737 };
1738 let (a, b) = (mk(1), mk(2));
1739 let s = sparse_sparse_matmul(a.as_ref(), b.as_ref(), 1.0, Par::Seq).unwrap();
1740 let p = sparse_sparse_matmul(a.as_ref(), b.as_ref(), 1.0, Par::rayon(4)).unwrap();
1741 close(s.val(), p.val(), 1e-12);
1742 let _ = Accum::Add;
1743 }
1744}