1use core::cell::UnsafeCell;
16use rayon::iter::plumbing::{Producer, ProducerCallback};
17use rayon::iter::{IndexedParallelIterator, IntoParallelIterator};
18use threadteam::TeamExec;
19
20pub fn with_team<R>(width: usize, f: impl FnOnce(crate::Par<'_>) -> R) -> R {
83 #[cfg(test)]
86 let _gate = team_turn();
87 let width = width.max(1);
88 if width == 1 {
89 return f(crate::Par::Seq);
90 }
91 let pool = threadteam::Pool::builder()
92 .threads(width)
93 .acquire_timeout(core::time::Duration::from_secs(30))
94 .build()
95 .expect("a thread pool for the team");
96 let team = pool
97 .into_team(width)
98 .expect("an idle pool: it was just created");
99 let mut exec = team.exec(width);
100 exec.team(|team| f(crate::Par::Team(team)))
101 .ok()
102 .expect("a team token for an idle team")
103}
104
105
106#[cfg(test)]
113pub(crate) static TEAM_GATE: std::sync::Mutex<()> = std::sync::Mutex::new(());
114
115#[cfg(test)]
117pub(crate) struct TeamTurn(Option<std::sync::MutexGuard<'static, ()>>);
118#[cfg(test)]
119thread_local! {
120 static TEAM_DEPTH: core::cell::Cell<usize> = const { core::cell::Cell::new(0) };
121}
122#[cfg(test)]
123impl Drop for TeamTurn {
124 fn drop(&mut self) {
125 TEAM_DEPTH.with(|d| d.set(d.get() - 1));
126 drop(self.0.take());
127 }
128}
129#[cfg(test)]
131pub(crate) fn team_turn() -> TeamTurn {
132 let nested = TEAM_DEPTH.with(|d| {
133 d.set(d.get() + 1);
134 d.get() > 1
135 });
136 if nested {
137 TeamTurn(None)
138 } else {
139 TeamTurn(Some(TEAM_GATE.lock().unwrap_or_else(|e| e.into_inner())))
140 }
141}
142
143pub mod stats {
149 use alloc::vec::Vec;
150 use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering::Relaxed};
151 pub(super) static ON: AtomicBool = AtomicBool::new(false);
152 const W: usize = 17;
153 #[allow(clippy::declare_interior_mutable_const)]
154 const Z: AtomicUsize = AtomicUsize::new(0);
155 pub(super) static RUN: [AtomicUsize; W] = [Z; W];
157 pub(super) static SPLIT: [AtomicUsize; W] = [Z; W];
159 const WW: usize = 257;
161 pub(super) static WIDTH: [AtomicUsize; WW] = [Z; WW];
163 pub(super) static GEMM_BELOW_FLOOR: AtomicUsize = AtomicUsize::new(0);
164 pub(super) static GEMM_ONE_TILE: AtomicUsize = AtomicUsize::new(0);
165 pub(super) static GEMM_TILED: AtomicUsize = AtomicUsize::new(0);
166 pub(super) static GEMM_TEAM_KERNEL: AtomicUsize = AtomicUsize::new(0);
167 pub(super) fn hist<const N: usize>(h: &[AtomicUsize; N], n: usize) {
168 if ON.load(Relaxed) {
169 h[n.min(N - 1)].fetch_add(1, Relaxed);
170 }
171 }
172 pub(super) fn bump(c: &AtomicUsize) {
173 if ON.load(Relaxed) {
174 c.fetch_add(1, Relaxed);
175 }
176 }
177 #[derive(Debug, Clone, Default)]
179 pub struct Stats {
180 pub run: Vec<usize>,
182 pub split: Vec<usize>,
184 pub width: Vec<usize>,
186 pub gemm_below_floor: usize,
188 pub gemm_one_tile: usize,
190 pub gemm_tiled: usize,
192 pub gemm_team_kernel: usize,
195 }
196 pub fn enable(on: bool) {
198 ON.store(on, Relaxed);
199 }
200 pub fn take() -> Stats {
202 fn t<const N: usize>(h: &[AtomicUsize; N]) -> Vec<usize> {
203 h.iter().map(|c| c.swap(0, Relaxed)).collect()
204 }
205 Stats {
206 run: t(&RUN),
207 split: t(&SPLIT),
208 width: t(&WIDTH),
209 gemm_below_floor: GEMM_BELOW_FLOOR.swap(0, Relaxed),
210 gemm_one_tile: GEMM_ONE_TILE.swap(0, Relaxed),
211 gemm_tiled: GEMM_TILED.swap(0, Relaxed),
212 gemm_team_kernel: GEMM_TEAM_KERNEL.swap(0, Relaxed),
213 }
214 }
215}
216
217pub(crate) fn note_width(width: usize) {
219 stats::hist(&stats::WIDTH, width);
220}
221
222pub(crate) fn note_split(parts: usize) {
224 stats::hist(&stats::SPLIT, parts);
225}
226
227pub(crate) fn run(team: TeamExec<'_>, n: usize, f: &(dyn Fn(usize, usize) + Sync)) {
230 let size = team.size().get();
231 let m = n.min(size);
232 if m <= 1 {
233 stats::hist(&stats::RUN, 1);
234 return f(0, 1);
235 }
236 stats::hist(&stats::RUN, m);
237 if m == size {
238 team.run(|i, k| f(i, k));
239 } else {
240 let _ = team.split(&[m], |_c, child| {
241 child.run(|i, k| f(i, k));
242 });
243 }
244}
245
246pub(crate) fn for_each<T: Send, I>(
249 team: TeamExec<'_>,
250 n_jobs: usize,
251 iter: I,
252 f: impl Sync + Fn(T),
253) where
254 I: IntoParallelIterator<Item = T>,
255 I::Iter: IndexedParallelIterator,
256{
257 for_each_child(team, n_jobs, iter, |item, _child| f(item))
258}
259
260pub(crate) fn for_each_child<T: Send, I>(
267 team: TeamExec<'_>,
268 n_jobs: usize,
269 iter: I,
270 f: impl Sync + for<'c> Fn(T, TeamExec<'c>),
271) where
272 I: IntoParallelIterator<Item = T>,
273 I::Iter: IndexedParallelIterator,
274{
275 let iter = iter.into_par_iter();
276 let len = iter.len();
277 if len == 0 {
278 return;
279 }
280 let jobs = n_jobs.max(1).min(len).min(team.size().get());
281 stats::hist(&stats::SPLIT, jobs);
282 struct Deal<'t, 'f, F> {
283 team: TeamExec<'t>,
284 len: usize,
285 jobs: usize,
286 f: &'f F,
287 }
288 impl<T, F: Sync + for<'c> Fn(T, TeamExec<'c>)> ProducerCallback<T> for Deal<'_, '_, F> {
289 type Output = ();
290 fn callback<P: Producer<Item = T>>(self, producer: P) {
291 deal(self.team, producer, self.len, self.jobs, self.f)
294 }
295 }
296 iter.with_producer(Deal { team, len, jobs, f: &f });
297}
298
299const FLAT_JOBS: usize = 64;
304const FLAT_BYTES: usize = 256;
305
306struct Piece<P>(UnsafeCell<Option<P>>);
308unsafe impl<P: Send> Sync for Piece<P> {}
310impl<P> Piece<P> {
311 unsafe fn take(&self) -> P {
314 unsafe { (*self.0.get()).take() }.expect("each piece is taken once")
315 }
316}
317
318fn deal<T, P: Producer<Item = T>, F: Sync + for<'c> Fn(T, TeamExec<'c>)>(
321 mut team: TeamExec<'_>,
322 producer: P,
323 len: usize,
324 jobs: usize,
325 f: &F,
326) {
327 if jobs <= 1 {
328 for item in producer.into_iter() {
329 f(item, team.reborrow());
330 }
331 return;
332 }
333 if jobs <= FLAT_JOBS && core::mem::size_of::<Option<P>>() <= FLAT_BYTES {
334 let mut pieces: [Piece<P>; FLAT_JOBS] = core::array::from_fn(|_| Piece(UnsafeCell::new(None)));
335 let mut rest = producer;
336 for (c, piece) in pieces.iter_mut().enumerate().take(jobs - 1) {
337 let (_, n) = super::thread::par_split_indices(len, c, jobs);
339 let (head, tail) = rest.split_at(n);
340 *piece.0.get_mut() = Some(head);
341 rest = tail;
342 }
343 *pieces[jobs - 1].0.get_mut() = Some(rest);
344 team.fork(jobs, |c, mut child| {
345 let piece = unsafe { pieces[c].take() };
347 for item in piece.into_iter() {
348 f(item, child.reborrow());
349 }
350 });
351 return;
352 }
353 let left = jobs / 2;
355 let (left_len, _) = super::thread::par_split_indices(len, left, jobs);
356 let size = team.size().get();
357 let (left_size, _) = super::thread::par_split_indices(size, left, jobs);
358 let (a, b) = producer.split_at(left_len);
359 let halves = [Piece(UnsafeCell::new(Some(a))), Piece(UnsafeCell::new(Some(b)))];
360 note_split(2);
361 team.split(&[left_size, size - left_size], |c, child| {
362 let half = unsafe { halves[c].take() };
364 if c == 0 {
365 deal(child, half, left_len, left, f)
366 } else {
367 deal(child, half, len - left_len, jobs - left, f)
368 }
369 })
370 .expect("two non-empty parts within the team");
371}
372
373struct SendPtr<T>(*mut T);
374unsafe impl<T> Send for SendPtr<T> {}
375unsafe impl<T> Sync for SendPtr<T> {}
376impl<T> Clone for SendPtr<T> {
377 fn clone(&self) -> Self {
378 *self
379 }
380}
381impl<T> Copy for SendPtr<T> {}
382impl<T> SendPtr<T> {
383 fn get(self) -> *mut T {
384 self.0
385 }
386}
387
388const ROW_GRAIN: usize = 8;
392const COL_GRAIN: usize = 4;
393
394pub(crate) fn tile_grid(m: usize, n: usize, nt: usize) -> (usize, usize) {
398 let max_pm = (m / ROW_GRAIN).max(1);
399 let max_pn = (n / COL_GRAIN).max(1);
400 let mut best = (1, 1);
401 let mut best_used = 1;
402 let mut best_cost = f64::INFINITY;
403 for pm in 1..=nt.min(max_pm) {
404 for pn in 1..=(nt / pm).min(max_pn) {
405 let used = pm * pn;
406 let cost = m as f64 / pm as f64 + n as f64 / pn as f64;
407 if used > best_used || (used == best_used && cost < best_cost) {
409 best = (pm, pn);
410 best_used = used;
411 best_cost = cost;
412 }
413 }
414 }
415 best
416}
417
418fn grain_range(len: usize, parts: usize, idx: usize, grain: usize) -> (usize, usize) {
421 let units = len.div_ceil(grain);
422 let edge = |i: usize| -> usize {
423 let base = units / parts;
424 let rem = units % parts;
425 let u = if i < rem { i * (base + 1) } else { rem + i * base };
426 (u * grain).min(len)
427 };
428 (edge(idx), edge(idx + 1))
429}
430
431#[cfg(all(target_arch = "x86_64", feature = "std"))]
433fn x86_kernel<E: 'static>() -> Option<(private_gemm_x86::InstrSet, private_gemm_x86::DType)> {
434 use core::any::TypeId;
435 use private_gemm_x86::*;
436 let id = TypeId::of::<E>();
437 let dtype = if id == TypeId::of::<f64>() {
438 DType::F64
439 } else if id == TypeId::of::<f32>() {
440 DType::F32
441 } else if id == TypeId::of::<num_complex::Complex<f64>>() {
442 DType::C64
443 } else if id == TypeId::of::<num_complex::Complex<f32>>() {
444 DType::C32
445 } else {
446 return None;
447 };
448 if std::arch::is_x86_feature_detected!("avx512f") {
449 Some((InstrSet::Avx512, dtype))
450 } else if std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("fma") {
451 Some((InstrSet::Avx256, dtype))
452 } else {
453 None
454 }
455}
456
457#[cfg(all(target_arch = "x86_64", feature = "std"))]
462pub(crate) mod kernel {
463 use super::{stats, x86_kernel};
464 use core::cell::{RefCell, UnsafeCell};
465 use core::time::Duration;
466 use private_gemm_x86::Fork;
467 use threadteam::TeamExec;
468
469 struct TeamFork<'t> {
474 team: RefCell<TeamExec<'t>>,
475 width: usize,
476 }
477
478 impl Fork for TeamFork<'_> {
479 fn width(&self) -> usize {
480 self.width
481 }
482
483 fn run(&self, f: &(dyn Fn(usize) + Sync)) {
484 stats::hist(&stats::RUN, self.width);
485 self.team.borrow_mut().reborrow().run(|j, _| f(j));
487 }
488 }
489
490 pub(crate) fn kernel_width(team: &TeamExec<'_>, estimate: Duration) -> Option<usize> {
494 match crate::utils::cost::par_width_team(team, estimate, 2, usize::MAX) {
496 1 => {
497 stats::bump(&stats::GEMM_BELOW_FLOOR);
498 None
499 },
500 width => Some(width),
501 }
502 }
503
504 pub(crate) fn with_fork<R>(team: TeamExec<'_>, width: usize, f: impl FnOnce(&dyn Fork) -> R) -> R {
507 stats::bump(&stats::GEMM_TEAM_KERNEL);
508 let size = team.size().get();
509 let width = width.clamp(1, size);
510 if width < size {
511 super::note_width(width);
512 }
513 if width == size {
514 return f(&TeamFork { team: RefCell::new(team), width });
515 }
516 struct Slot<F, R>(UnsafeCell<(Option<F>, Option<R>)>);
517 unsafe impl<F, R> Sync for Slot<F, R> {}
520 impl<F, R> Slot<F, R> {
521 fn get(&self) -> *mut (Option<F>, Option<R>) {
522 self.0.get()
523 }
524 }
525 let slot = Slot(UnsafeCell::new((Some(f), None)));
526 team.split(&[width], |_, child| {
527 let (f, out) = unsafe { &mut *slot.get() };
529 let f = f.take().expect("one child");
530 *out = Some(f(&TeamFork { team: RefCell::new(child), width }));
531 })
532 .expect("1 <= width <= size is a valid split");
533 slot.0.into_inner().1.expect("the child ran")
534 }
535
536 pub(crate) fn available<E: 'static>() -> Option<(private_gemm_x86::InstrSet, private_gemm_x86::DType)> {
538 x86_kernel::<E>()
539 }
540}
541
542#[allow(clippy::too_many_arguments)]
552pub(crate) unsafe fn gemm_team<E: 'static + Copy + Send + Sync>(
553 team: TeamExec<'_>,
554 m: usize,
555 n: usize,
556 k: usize,
557 ns_per_mac: f64,
558 dst: *mut E,
559 dst_cs: isize,
560 dst_rs: isize,
561 read_dst: bool,
562 lhs: *const E,
563 lhs_cs: isize,
564 lhs_rs: isize,
565 rhs: *const E,
566 rhs_cs: isize,
567 rhs_rs: isize,
568 alpha: E,
569 beta: E,
570 conj_lhs: bool,
571 conj_rhs: bool,
572) {
573 #[cfg(all(target_arch = "x86_64", feature = "std"))]
574 if let Some((feat, dtype)) = kernel::available::<E>() {
575 use private_gemm_x86::*;
576 macro_rules! go {
580 ($call:path, $last:expr) => {
581 $call(
582 dtype,
583 IType::U64,
584 feat,
585 m,
586 n,
587 k,
588 dst as *mut (),
589 dst_rs,
590 dst_cs,
591 core::ptr::null(),
592 core::ptr::null(),
593 DstKind::Full,
594 if read_dst { Accum::Add } else { Accum::Replace },
595 lhs as *const (),
596 lhs_rs,
597 lhs_cs,
598 conj_lhs,
599 core::ptr::null(),
600 0,
601 rhs as *const (),
602 rhs_rs,
603 rhs_cs,
604 conj_rhs,
605 &raw const beta as *const (),
606 $last,
607 )
608 };
609 }
610 let call = |fork: Option<&dyn Fork>| unsafe {
611 match fork {
612 Some(fork) => go!(gemm_with, fork),
613 None => go!(gemm, 1),
614 }
615 };
616 let estimate = crate::utils::cost::serial_time(m.saturating_mul(n).saturating_mul(k), ns_per_mac);
617 match kernel::kernel_width(&team, estimate) {
618 Some(width) => kernel::with_fork(team, width, |fork| call(Some(fork))),
619 None => call(None),
620 }
621 return;
622 }
623 gemm_tiled::<E>(team, m, n, k, ns_per_mac, dst, dst_cs, dst_rs, read_dst, lhs, lhs_cs, lhs_rs, rhs, rhs_cs, rhs_rs, alpha, beta, conj_lhs, conj_rhs)
624}
625
626#[allow(clippy::too_many_arguments)]
640pub(crate) unsafe fn gemm_tiled<E: 'static + Copy + Send + Sync>(
641 team: TeamExec<'_>,
642 m: usize,
643 n: usize,
644 k: usize,
645 ns_per_mac: f64,
646 dst: *mut E,
647 dst_cs: isize,
648 dst_rs: isize,
649 read_dst: bool,
650 lhs: *const E,
651 lhs_cs: isize,
652 lhs_rs: isize,
653 rhs: *const E,
654 rhs_cs: isize,
655 rhs_rs: isize,
656 alpha: E,
657 beta: E,
658 conj_lhs: bool,
659 conj_rhs: bool,
660) {
661 let seq = |dst: *mut E, lhs: *const E, rhs: *const E, m: usize, n: usize| unsafe {
665 gemm::gemm(
666 m,
667 n,
668 k,
669 dst,
670 dst_cs,
671 dst_rs,
672 read_dst,
673 lhs,
674 lhs_cs,
675 lhs_rs,
676 rhs,
677 rhs_cs,
678 rhs_rs,
679 alpha,
680 beta,
681 false,
682 conj_lhs,
683 conj_rhs,
684 gemm::Parallelism::None,
685 )
686 };
687 let estimate = crate::utils::cost::serial_time(m.saturating_mul(n).saturating_mul(k), ns_per_mac);
691 let nt = crate::utils::cost::par_width_team(&team, estimate, 1, usize::MAX);
692 if nt <= 1 {
693 stats::bump(&stats::GEMM_BELOW_FLOOR);
694 return seq(dst, lhs, rhs, m, n);
695 }
696 let (pm, pn) = tile_grid(m, n, nt);
697 if pm * pn <= 1 {
698 stats::bump(&stats::GEMM_ONE_TILE);
699 return seq(dst, lhs, rhs, m, n);
700 }
701 if pm * pn < nt
704 && !crate::utils::cost::pays_at(estimate, 1, team.dispatch_cost(core::num::NonZeroUsize::new(pm * pn).unwrap()))
705 {
706 stats::bump(&stats::GEMM_BELOW_FLOOR);
707 return seq(dst, lhs, rhs, m, n);
708 }
709 if pm * pn < team.size().get() {
710 note_width(pm * pn);
711 }
712 stats::bump(&stats::GEMM_TILED);
713 let (dst, lhs, rhs) = (
714 SendPtr(dst),
715 SendPtr(lhs as *mut E),
716 SendPtr(rhs as *mut E),
717 );
718 run(team, pm * pn, &|i, parts| {
719 let mut t = i;
721 while t < pm * pn {
722 let (ti, tj) = (t % pm, t / pm);
723 let (i0, i1) = grain_range(m, pm, ti, ROW_GRAIN);
724 let (j0, j1) = grain_range(n, pn, tj, COL_GRAIN);
725 if i1 > i0 && j1 > j0 {
726 let d = dst.get().wrapping_offset(i0 as isize * dst_rs + j0 as isize * dst_cs);
727 let l = lhs.get().wrapping_offset(i0 as isize * lhs_rs) as *const E;
728 let r = rhs.get().wrapping_offset(j0 as isize * rhs_cs) as *const E;
729 seq(d, l, r, i1 - i0, j1 - j0);
730 }
731 t += parts;
732 }
733 });
734}
735
736#[cfg(test)]
737mod tests {
738 use super::*;
739 use crate::linalg::matmul::matmul;
740 use crate::{c64, Accum, Mat, MatRef, Par};
741 use core::sync::atomic::{AtomicUsize, Ordering};
742 use threadteam::Pool;
743 use core::time::Duration as StdDuration;
744
745 fn team_of(width: usize) -> threadteam::Team {
746 Pool::builder()
747 .threads(width)
748 .acquire_timeout(std::time::Duration::from_secs(30))
749 .build()
750 .unwrap()
751 .into_team(width)
752 .expect("an idle pool")
753 }
754
755 use super::team_turn;
757
758 fn with_par<R>(width: usize, f: impl FnOnce(Par<'_>) -> R) -> R {
760 let _turn = team_turn();
761 let team = team_of(width);
762 let mut exec = team.exec(width);
763 exec.team(|team| f(Par::Team(team))).ok().expect("a team token")
764 }
765
766 #[test]
770 fn with_team_has_width_participants() {
771 for width in [2usize, 3, 4] {
772 let size = with_team(width, |par| match par {
773 Par::Team(t) => t.size().get(),
774 _ => 1,
775 });
776 assert!(size == width);
777 }
778 let pool = Pool::builder()
779 .threads(3)
780 .acquire_timeout(std::time::Duration::from_secs(30))
781 .build()
782 .unwrap();
783 let team = pool.into_team(4).ok().expect("an idle pool");
784 assert!(team.size().get() == 3);
785 }
786
787 fn rand_mat(m: usize, n: usize, seed: u64) -> Mat<f64> {
788 let mut x = seed.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
789 Mat::from_fn(m, n, |_, _| {
790 x = x.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
791 ((x >> 11) as f64 / (1u64 << 53) as f64) - 0.5
792 })
793 }
794
795 fn max_diff(a: MatRef<'_, f64>, b: MatRef<'_, f64>) -> f64 {
796 let mut d: f64 = 0.0;
797 for j in 0..a.ncols() {
798 for i in 0..a.nrows() {
799 d = d.max((a[(i, j)] - b[(i, j)]).abs());
800 }
801 }
802 d
803 }
804
805 #[test]
810 fn small_qr_on_a_team_is_the_sequential_one() {
811 use crate::dyn_stack::{MemBuffer, MemStack};
812 use crate::linalg::qr::no_pivoting::factor as qrf;
813 let n = 32;
814 let a0 = rand_mat(n, n, 11);
815 let run = |mut par: Par<'_>| {
816 let bs = qrf::recommended_block_size::<f64>(n, n);
817 let mut a = a0.clone();
818 let mut q = Mat::<f64>::zeros(bs, n);
819 let mut mem = MemBuffer::new(qrf::qr_in_place_scratch::<f64>(n, n, bs, par.shape(), Default::default()));
820 qrf::qr_in_place(a.as_mut(), q.as_mut(), par.reborrow(), MemStack::new(&mut mem), Default::default());
821 (a, q)
822 };
823 let (a_seq, q_seq) = run(Par::Seq);
824 let (a_team, q_team) = with_par(4, |par| {
825 stats::take();
826 stats::enable(true);
827 let r = run(par);
828 stats::enable(false);
829 let s = stats::take();
830 let sections: usize = s.run.iter().skip(2).sum::<usize>() + s.split.iter().skip(2).sum::<usize>();
831 assert_eq!(sections, 0, "a QR {n} opened a section: {s:?}");
832 assert_eq!(s.gemm_tiled + s.gemm_team_kernel, 0, "{s:?}");
833 r
834 });
835 assert_eq!(max_diff(a_seq.as_ref(), a_team.as_ref()), 0.0);
836 assert_eq!(max_diff(q_seq.as_ref(), q_team.as_ref()), 0.0);
837 }
838
839 #[test]
843 #[cfg(all(target_arch = "x86_64", feature = "std"))]
844 fn kernel_width_rejects_at_the_cheapest_width() {
845 use threadteam::{DispatchCosts, DispatchKind, WidthCosts};
846 use core::time::Duration;
847 let _turn = team_turn();
848 let us = Duration::from_micros;
849 let costs = DispatchCosts::for_class(threadteam::HostClass::Generic).with(DispatchKind::Team, WidthCosts::new([us(2), us(7), us(12)]));
850 let team = Pool::builder()
851 .threads(6)
852 .acquire_timeout(std::time::Duration::from_secs(30))
853 .dispatch_costs(costs)
854 .build()
855 .unwrap()
856 .into_team(6)
857 .expect("an idle pool");
858 let mut exec = team.exec(6);
859 exec.team(|team| {
860 assert_eq!(kernel::kernel_width(&team, us(13)), Some(2));
861 assert_eq!(kernel::kernel_width(&team, us(11)), None);
862 let par = Par::Team(team);
863 assert!(crate::utils::cost::par_can_pay_time(&par, us(6), 1));
864 assert!(!crate::utils::cost::par_can_pay_time(&par, us(5), 1));
865 })
866 .ok()
867 .expect("a team token");
868 }
869
870 #[test]
874 #[cfg(all(target_arch = "x86_64", feature = "std"))]
875 fn kernel_width_finds_the_widest_paying_width_on_a_falling_row() {
876 use threadteam::{DispatchCosts, DispatchKind, WidthCosts};
877 use core::time::Duration;
878 let _turn = team_turn();
879 let us = Duration::from_micros;
880 let costs = DispatchCosts::for_class(threadteam::HostClass::Generic).with(DispatchKind::Team, WidthCosts::new([us(3), us(2), us(1)]));
881 let team = Pool::builder()
882 .threads(6)
883 .acquire_timeout(std::time::Duration::from_secs(30))
884 .dispatch_costs(costs)
885 .build()
886 .unwrap()
887 .into_team(6)
888 .expect("an idle pool");
889 let mut exec = team.exec(6);
890 exec.team(|team| {
891 assert_eq!(kernel::kernel_width(&team, us(7)), Some(6));
892 assert_eq!(kernel::kernel_width(&team, us(5)), None);
893 })
894 .ok()
895 .expect("a team token");
896 }
897
898 fn with_row<R>(width: usize, anchors: &[(usize, StdDuration)], f: impl FnOnce(Par<'_>) -> R) -> R {
900 use threadteam::{DispatchCosts, DispatchKind, WidthCosts};
901 let _turn = team_turn();
902 let costs = DispatchCosts::for_class(threadteam::HostClass::Generic)
903 .with(DispatchKind::Team, WidthCosts::from_anchors(anchors).expect("a valid row"));
904 let team = Pool::builder()
905 .threads(width)
906 .acquire_timeout(std::time::Duration::from_secs(30))
907 .dispatch_costs(costs)
908 .build()
909 .unwrap()
910 .into_team(width)
911 .expect("an idle pool");
912 let mut exec = team.exec(width);
913 exec.team(|team| f(Par::Team(team))).ok().expect("a team token")
914 }
915
916 #[test]
920 fn par_width_takes_the_widest_width_that_pays() {
921 use crate::utils::cost::{par_width_time, par_width_time_upto};
922 let us = StdDuration::from_micros;
923 with_row(8, &[(2, us(1)), (4, us(2)), (8, us(100))], |par| {
926 assert_eq!(par_width_time(&par, us(6), 1), 4);
927 assert_eq!(par_width_time(&par, us(5), 1), 3);
929 assert_eq!(par_width_time(&par, us(2), 1), 1);
930 assert_eq!(par_width_time(&par, us(2000), 1), 8);
931 });
932 with_row(8, &[(2, us(3)), (4, us(9)), (6, us(1)), (8, us(1))], |par| {
934 assert_eq!(par_width_time(&par, us(3), 1), 8);
935 assert_eq!(par_width_time_upto(&par, us(3), 1, 5), 1);
936 assert_eq!(par_width_time_upto(&par, us(9), 1, 5), 2);
937 });
938 assert_eq!(par_width_time(&Par::Seq, us(1_000_000), 1), 1);
939 }
940
941 #[test]
944 fn with_width_runs_on_a_sub_team_of_that_width() {
945 use crate::utils::thread::with_width;
946 with_par(8, |mut par| {
947 stats::enable(true);
948 let _ = stats::take();
949 for (w, want) in [(3, 3), (8, 8), (9, 8), (1, 1), (0, 1)] {
950 let seen = AtomicUsize::new(0);
951 let degree = with_width(par.reborrow(), w, |par| {
952 let degree = par.degree();
953 if let Par::Team(team) = par {
954 team.run(|_, _| {
955 seen.fetch_add(1, Ordering::Relaxed);
956 });
957 }
958 degree
959 });
960 assert_eq!(degree, want, "width {w}");
961 if want > 1 {
962 assert_eq!(seen.load(Ordering::Relaxed), want, "width {w}");
963 }
964 }
965 assert_eq!(par.degree(), 8);
966 stats::enable(false);
967 let s = stats::take();
968 assert_eq!(s.width[3], 1, "{s:?}");
970 assert_eq!(s.width.iter().sum::<usize>(), 1, "{s:?}");
971 });
972 }
973
974 #[test]
977 fn the_families_narrow_their_regions_and_match_sequential() {
978 use crate::dyn_stack::{MemBuffer, MemStack};
979 use crate::linalg;
980 let ns = StdDuration::from_nanos;
981 let row = [(2, ns(10)), (4, ns(10)), (5, StdDuration::from_secs(1))];
983 let tol = |a: MatRef<'_, f64>, b: MatRef<'_, f64>, what: &str| {
984 let scale = (0..a.ncols()).flat_map(|j| (0..a.nrows()).map(move |i| (i, j))).fold(1.0f64, |m, (i, j)| m.max(a[(i, j)].abs()));
985 let d = max_diff(a, b);
986 assert!(d <= 1e-10 * scale, "{what}: {d} against {scale}");
987 };
988 type Run = fn(Par<'_>) -> Vec<Mat<f64>>;
989 let families: [(&str, Run); 14] = [
990 ("full pivoting lu", |par| {
991 let n = 200;
992 let mut a = rand_mat(n, n, 11);
993 let (mut p, mut pi, mut q, mut qi) = (vec![0usize; n], vec![0usize; n], vec![0usize; n], vec![0usize; n]);
994 let req = linalg::lu::full_pivoting::factor::lu_in_place_scratch::<usize, f64>(n, n, par.shape(), Default::default());
995 linalg::lu::full_pivoting::factor::lu_in_place(a.as_mut(), &mut p, &mut pi, &mut q, &mut qi, par, MemStack::new(&mut MemBuffer::new(req)), Default::default());
996 vec![a]
997 }),
998 ("partial pivoting lu", |par| {
999 let n = 400;
1000 let mut a = rand_mat(n, n, 12);
1001 let (mut p, mut pi) = (vec![0usize; n], vec![0usize; n]);
1002 let req = linalg::lu::partial_pivoting::factor::lu_in_place_scratch::<usize, f64>(n, n, par.shape(), Default::default());
1003 linalg::lu::partial_pivoting::factor::lu_in_place(a.as_mut(), &mut p, &mut pi, par, MemStack::new(&mut MemBuffer::new(req)), Default::default());
1004 vec![a]
1005 }),
1006 ("column pivoting qr", |par| {
1007 let n = 200;
1008 let mut a = rand_mat(n, n, 13);
1009 let bs = linalg::qr::no_pivoting::factor::recommended_block_size::<f64>(n, n);
1010 let mut h = Mat::<f64>::zeros(bs, n);
1011 let (mut q, mut qi) = (vec![0usize; n], vec![0usize; n]);
1012 let req = linalg::qr::col_pivoting::factor::qr_in_place_scratch::<usize, f64>(n, n, bs, par.shape(), Default::default());
1013 linalg::qr::col_pivoting::factor::qr_in_place(a.as_mut(), h.as_mut(), &mut q, &mut qi, par, MemStack::new(&mut MemBuffer::new(req)), Default::default());
1014 vec![a, h]
1015 }),
1016 ("hessenberg", |par| {
1017 let n = 160;
1018 let mut a = rand_mat(n, n, 14);
1019 let bs = 8;
1020 let mut h = Mat::<f64>::zeros(bs, n - 1);
1021 let req = linalg::evd::hessenberg::hessenberg_in_place_scratch::<f64>(n, bs, par.shape(), Default::default());
1022 linalg::evd::hessenberg::hessenberg_in_place(a.as_mut(), h.as_mut(), par, MemStack::new(&mut MemBuffer::new(req)), Default::default());
1023 vec![a, h]
1024 }),
1025 ("tridiagonalization", |par| {
1026 let n = 200;
1027 let x = rand_mat(n, n, 15);
1028 let mut a = &x + x.transpose();
1029 let bs = 8;
1030 let mut h = Mat::<f64>::zeros(bs, n - 1);
1031 let req = linalg::evd::tridiag::tridiag_in_place_scratch::<f64>(n, par.shape(), Default::default());
1032 linalg::evd::tridiag::tridiag_in_place(a.as_mut(), h.as_mut(), par, MemStack::new(&mut MemBuffer::new(req)), Default::default());
1033 vec![a, h]
1034 }),
1035 ("bidiagonalization", |par| {
1036 let (m, n) = (300, 200);
1037 let mut a = rand_mat(m, n, 16);
1038 let bs = 8;
1039 let mut hl = Mat::<f64>::zeros(bs, n);
1040 let mut hr = Mat::<f64>::zeros(bs, n - 1);
1041 let req = linalg::svd::bidiag::bidiag_in_place_scratch::<f64>(m, n, par.shape(), Default::default());
1042 linalg::svd::bidiag::bidiag_in_place(a.as_mut(), hl.as_mut(), hr.as_mut(), par, MemStack::new(&mut MemBuffer::new(req)), Default::default());
1043 vec![a, hl, hr]
1044 }),
1045 ("bunch kaufman, full pivoting", |par| {
1046 use linalg::cholesky::lblt::factor::{LbltParams, PivotingStrategy};
1047 let n = 200;
1048 let x = rand_mat(n, n, 17);
1049 let mut a = &x + x.transpose();
1050 let mut sub = crate::diag::Diag::<f64>::zeros(n);
1051 let (mut p, mut pi) = (vec![0usize; n], vec![0usize; n]);
1052 let params = LbltParams {
1053 pivoting: PivotingStrategy::Full,
1054 ..<LbltParams as crate::Auto<f64>>::auto()
1055 };
1056 let req = linalg::cholesky::lblt::factor::cholesky_in_place_scratch::<usize, f64>(n, par.shape(), params.into());
1057 linalg::cholesky::lblt::factor::cholesky_in_place(a.as_mut(), sub.as_mut(), &mut p, &mut pi, par, MemStack::new(&mut MemBuffer::new(req)), params.into());
1058 vec![a, sub.column_vector().as_mat().to_owned()]
1059 }),
1060 ("triangular solve", |par| {
1061 let (n, k) = (100, 512);
1062 let l = rand_mat(n, n, 18) + Mat::<f64>::identity(n, n) * (n as f64);
1063 let mut x = rand_mat(n, k, 19);
1064 linalg::triangular_solve::solve_lower_triangular_in_place(l.as_ref(), x.as_mut(), par);
1065 vec![x]
1066 }),
1067 ("matvec, row split", |par| {
1068 let a = rand_mat(1000, 1000, 20);
1069 let v = rand_mat(1000, 1, 21);
1070 let mut y = Mat::<f64>::zeros(1000, 1);
1071 matmul(y.as_mut(), Accum::Replace, a.as_ref(), v.as_ref(), 1.0, par);
1072 vec![y]
1073 }),
1074 ("matvec, column split", |par| {
1075 let a = rand_mat(1000, 1000, 22);
1076 let v = rand_mat(1000, 1, 23);
1077 let mut y = Mat::<f64>::zeros(1, 1000);
1078 matmul(y.as_mut(), Accum::Replace, v.transpose(), a.as_ref(), 1.0, par);
1079 vec![y]
1080 }),
1081 ("block householder apply", |par| {
1082 let (m, b, k) = (2000, 8, 256);
1083 let basis = rand_mat(m, b, 24);
1084 let factor = Mat::<f64>::from_fn(b, b, |i, j| if i == j { 1.5 } else if i < j { 0.1 } else { 0.0 });
1085 let mut target = rand_mat(m, k, 25);
1086 let req = linalg::householder::apply_block_householder_on_the_left_in_place_scratch::<f64>(m, b, k);
1087 linalg::householder::apply_block_householder_on_the_left_in_place_with_conj(
1088 basis.as_ref(), factor.as_ref(), crate::Conj::No, target.as_mut(), par, MemStack::new(&mut MemBuffer::new(req)),
1089 );
1090 vec![target]
1091 }),
1092 ("householder factor upgrade", |par| {
1093 let (rows, b) = (4000, 16);
1094 let essentials = rand_mat(rows, b, 26);
1095 let mut factor = Mat::<f64>::from_fn(b, b, |i, j| if i == j { 1.25 } else { 0.0 });
1096 linalg::householder::upgrade_householder_factor(factor.as_mut(), essentials.as_ref(), b, 1, par);
1097 vec![factor]
1098 }),
1099 ("sparse product", |par| {
1100 use crate::sparse::{SparseColMat, Triplet};
1101 let n = 3000;
1102 let triplets: Vec<_> = (0..n)
1103 .flat_map(|j| (0..4).map(move |t| Triplet::new((j * 7 + t * 131) % n, j, ((j + t) as f64).sin())))
1104 .collect();
1105 let a = SparseColMat::<usize, f64>::try_new_from_triplets(n, n, &triplets).unwrap();
1106 let c = crate::sparse::linalg::matmul::sparse_sparse_matmul(a.as_ref(), a.as_ref(), 1.0, par).unwrap();
1107 vec![c.to_dense()]
1108 }),
1109 ("tiled product (the non-x86 route)", |par| {
1110 let (m, n, k) = (301, 263, 190);
1111 let a = rand_mat(m, k, 27);
1112 let b = rand_mat(k, n, 28);
1113 let mut c = Mat::<f64>::zeros(m, n);
1114 match par {
1115 Par::Team(team) => unsafe {
1116 gemm_tiled::<f64>(
1117 team, m, n, k, 1.0, c.as_ptr_mut(), c.col_stride(), c.row_stride(), false,
1118 a.as_ptr(), a.col_stride(), a.row_stride(), b.as_ptr(), b.col_stride(), b.row_stride(),
1119 0.0, 1.0, false, false,
1120 )
1121 },
1122 Par::Seq => matmul(c.as_mut(), Accum::Replace, a.as_ref(), b.as_ref(), 1.0, Par::Seq),
1123 }
1124 vec![c]
1125 }),
1126 ];
1127 for (name, run) in families {
1128 let want = run(Par::Seq);
1129 let (got, s) = with_row(8, &row, |par| {
1130 stats::enable(true);
1131 let _ = stats::take();
1132 let got = run(par);
1133 stats::enable(false);
1134 (got, stats::take())
1135 });
1136 assert!(s.width[4] > 0, "{name}: no region narrowed to 4 ({s:?})");
1138 assert!(s.split[4] + s.run[4] > 0, "{name}: no split or section of 4 ({s:?})");
1141 assert!(s.split.iter().skip(5).all(|&c| c == 0), "{name}: a split wider than 4 ({s:?})");
1142 for (k, (a, b)) in want.iter().zip(got.iter()).enumerate() {
1143 tol(a.as_ref(), b.as_ref(), &alloc::format!("{name} output {k}"));
1144 }
1145 }
1146 }
1147
1148 #[test]
1152 fn never_reaches_the_compact_factor_upgrades() {
1153 use crate::dyn_stack::{MemBuffer, MemStack};
1154 use crate::linalg;
1155 use crate::utils::cost::ParThreshold;
1156 let ns = StdDuration::from_nanos;
1157 let row = [(2, ns(10)), (4, ns(10)), (5, StdDuration::from_secs(1))];
1158 let s = with_row(8, &row, |mut par| {
1159 stats::enable(true);
1160 let _ = stats::take();
1161 {
1162 use linalg::qr::col_pivoting::factor::*;
1163 let (m, n) = (33, 9);
1164 let mut a = rand_mat(m, n, 31);
1165 let mut h = Mat::<f64>::zeros(n, n);
1166 let (mut q, mut qi) = (vec![0usize; n], vec![0usize; n]);
1167 let params = ColPivQrParams { par_threshold: ParThreshold::Never, ..<ColPivQrParams as crate::Auto<f64>>::auto() };
1168 let req = qr_in_place_scratch::<usize, f64>(m, n, n, par.shape(), params.into());
1169 qr_in_place(a.as_mut(), h.as_mut(), &mut q, &mut qi, par.reborrow(), MemStack::new(&mut MemBuffer::new(req)), params.into());
1170 }
1171 {
1172 use linalg::svd::bidiag::*;
1173 let (m, n) = (300, 200);
1174 let mut a = rand_mat(m, n, 32);
1175 let mut hl = Mat::<f64>::zeros(8, n);
1176 let mut hr = Mat::<f64>::zeros(8, n - 1);
1177 let params = BidiagParams { par_threshold: ParThreshold::Never, ..<BidiagParams as crate::Auto<f64>>::auto() };
1178 let req = bidiag_in_place_scratch::<f64>(m, n, par.shape(), params.into());
1179 bidiag_in_place(a.as_mut(), hl.as_mut(), hr.as_mut(), par, MemStack::new(&mut MemBuffer::new(req)), params.into());
1180 }
1181 stats::enable(false);
1182 stats::take()
1183 });
1184 let parallel = s.run.iter().skip(2).sum::<usize>() + s.split.iter().skip(2).sum::<usize>() + s.gemm_tiled + s.gemm_team_kernel;
1185 assert_eq!(parallel, 0, "{s:?}");
1186 }
1187
1188 #[test]
1193 fn for_each_child_deals_every_item_once() {
1194 use rayon::prelude::*;
1195 for (width, len, n_jobs, big) in [
1196 (8usize, 1000usize, 7usize, false),
1197 (8, 1000, 7, true),
1198 (70, 1000, 70, false),
1199 (70, 333, 69, true),
1200 (5, 3, 5, false),
1201 ] {
1202 with_par(width, |par| {
1203 let Par::Team(team) = par else { unreachable!() };
1204 let seen: Vec<AtomicUsize> = (0..len).map(|_| AtomicUsize::new(0)).collect();
1205 let widths: Vec<AtomicUsize> = (0..len).map(|_| AtomicUsize::new(0)).collect();
1206 let f = |i: usize, child: TeamExec<'_>| {
1207 seen[i].fetch_add(1, Ordering::Relaxed);
1208 widths[i].store(child.size().get(), Ordering::Relaxed);
1209 };
1210 if big {
1211 core::assert!(core::mem::size_of::<[u8; 512]>() > FLAT_BYTES);
1212 for_each_child(team, n_jobs, (0..len).into_par_iter().map_with([0u8; 512], |_, i| i), f);
1213 } else {
1214 for_each_child(team, n_jobs, (0..len).into_par_iter(), f);
1215 }
1216 let jobs = n_jobs.min(len).min(width);
1217 let (lo, hi) = (width / jobs, width.div_ceil(jobs));
1218 for i in 0..len {
1219 core::assert_eq!(seen[i].load(Ordering::Relaxed), 1, "{width} {len} {n_jobs} {big}: item {i}");
1220 let w = widths[i].load(Ordering::Relaxed);
1221 core::assert!(lo <= w && w <= hi, "{width} {len} {n_jobs} {big}: item {i} on a child of {w}");
1222 }
1223 });
1224 }
1225 }
1226
1227 #[test]
1230 fn for_each_child_unwinds_and_keeps_the_team() {
1231 use rayon::prelude::*;
1232 use std::panic::{catch_unwind, AssertUnwindSafe};
1233 struct Counted<'a>(&'a AtomicUsize);
1234 impl Drop for Counted<'_> {
1235 fn drop(&mut self) {
1236 self.0.fetch_add(1, Ordering::Relaxed);
1237 }
1238 }
1239 for width in [8usize, 70] {
1240 with_par(width, |mut par| {
1241 let dropped = AtomicUsize::new(0);
1242 let made = AtomicUsize::new(0);
1243 let r = catch_unwind(AssertUnwindSafe(|| {
1244 let Par::Team(team) = par.reborrow() else { unreachable!() };
1245 for_each_child(
1246 team,
1247 width,
1248 (0..1000usize).into_par_iter().map(|i| {
1249 made.fetch_add(1, Ordering::Relaxed);
1250 (i, Counted(&dropped))
1251 }),
1252 |(i, _c), _| {
1253 if i == 500 {
1254 panic!("block failure");
1255 }
1256 },
1257 );
1258 }));
1259 core::assert!(r.is_err());
1260 core::assert_eq!(dropped.load(Ordering::Relaxed), made.load(Ordering::Relaxed));
1261 let Par::Team(team) = par.reborrow() else { unreachable!() };
1263 let hits = AtomicUsize::new(0);
1264 for_each_child(team, width, (0..100usize).into_par_iter(), |_, _| {
1265 hits.fetch_add(1, Ordering::Relaxed);
1266 });
1267 core::assert_eq!(hits.load(Ordering::Relaxed), 100);
1268 });
1269 }
1270 }
1271
1272 #[test]
1274 fn par_pays_time_is_the_shape_gate() {
1275 use crate::utils::cost::par_pays_time;
1276 use core::time::Duration;
1277 with_par(6, |par| {
1278 let shape = par.shape();
1279 for &us in &[0u64, 1, 3, 5, 8, 12, 30, 1000] {
1280 for sections in 1..=3 {
1281 for width in 1..=8usize {
1282 let est = Duration::from_micros(us);
1283 let w = core::num::NonZeroUsize::new(width.min(6)).unwrap();
1284 let want = width >= 2 && shape.pays(est, sections, w);
1285 assert_eq!(par_pays_time(&par, est, sections, width), want, "{us} us, {sections}, {width}");
1286 }
1287 }
1288 }
1289 });
1290 assert!(!par_pays_time(&Par::Seq, Duration::from_secs(1), 1, 4));
1291 }
1292
1293 #[test]
1294 fn grid_uses_the_team_and_stays_in_bounds() {
1295 for nt in 1..=12 {
1296 for &(m, n) in &[(2048, 2048), (4096, 64), (64, 4096), (100, 7), (8, 4), (1, 1000)] {
1297 let (pm, pn) = tile_grid(m, n, nt);
1298 assert!(pm * pn <= nt && pm >= 1 && pn >= 1, "{m}x{n} nt={nt}");
1299 let mut rows = 0;
1301 for i in 0..pm {
1302 let (a, b) = grain_range(m, pm, i, ROW_GRAIN);
1303 assert_eq!(a, rows);
1304 rows = b;
1305 }
1306 assert_eq!(rows, m);
1307 let mut cols = 0;
1308 for j in 0..pn {
1309 let (a, b) = grain_range(n, pn, j, COL_GRAIN);
1310 assert_eq!(a, cols);
1311 cols = b;
1312 }
1313 assert_eq!(cols, n);
1314 }
1315 }
1316 assert_eq!(tile_grid(2048, 2048, 6).0 * tile_grid(2048, 2048, 6).1, 6);
1317 }
1318
1319 #[test]
1320 fn gemm_on_the_team_matches_sequential() {
1321 with_par(6, |mut par| {
1322 for &(m, n, k) in &[(512, 512, 512), (257, 301, 129), (4096, 64, 64), (64, 4096, 64), (1000, 33, 500)] {
1323 let a = rand_mat(m, k, 1);
1324 let b = rand_mat(k, n, 2);
1325 for beta in [Accum::Replace, Accum::Add] {
1326 let mut want = rand_mat(m, n, 3);
1327 let mut got = want.clone();
1328 matmul(want.as_mut(), beta, a.as_ref(), b.as_ref(), 1.5, Par::Seq);
1329 matmul(got.as_mut(), beta, a.as_ref(), b.as_ref(), 1.5, par.reborrow());
1330 assert!(max_diff(want.as_ref(), got.as_ref()) < 1e-10, "{m}x{n}x{k}");
1331 }
1332 let mut want = Mat::<f64>::zeros(m, n);
1334 let mut got = Mat::<f64>::zeros(m, n);
1335 let at = a.transpose().to_owned();
1336 matmul(want.as_mut(), Accum::Replace, at.transpose(), b.as_ref(), 1.0, Par::Seq);
1337 matmul(got.as_mut(), Accum::Replace, at.transpose(), b.as_ref(), 1.0, par.reborrow());
1338 assert!(max_diff(want.as_ref(), got.as_ref()) < 1e-10);
1339 }
1340 let (m, n, k) = (300, 280, 200);
1342 let a = Mat::<c64>::from_fn(m, k, |i, j| c64::new(i as f64 * 0.01 - j as f64 * 0.003, (i + j) as f64 * 0.002));
1343 let b = Mat::<c64>::from_fn(k, n, |i, j| c64::new((i * j % 7) as f64 * 0.1, i as f64 * 0.01 - j as f64 * 0.01));
1344 let mut want = Mat::<c64>::zeros(m, n);
1345 let mut got = Mat::<c64>::zeros(m, n);
1346 matmul(want.as_mut(), Accum::Replace, a.as_ref(), b.as_ref(), c64::new(1.0, 0.0), Par::Seq);
1347 matmul(got.as_mut(), Accum::Replace, a.as_ref(), b.as_ref(), c64::new(1.0, 0.0), par.reborrow());
1348 let mut d: f64 = 0.0;
1349 for j in 0..n {
1350 for i in 0..m {
1351 d = d.max((want[(i, j)] - got[(i, j)]).norm_sqr().sqrt());
1352 }
1353 }
1354 assert!(d < 1e-9);
1355 });
1356 }
1357
1358 #[test]
1362 fn the_team_floor_is_per_dtype() {
1363 let engaged = |st: &stats::Stats| st.gemm_tiled + st.gemm_team_kernel;
1365 let before = |m: usize, n: usize, k: usize, par: &mut Par<'_>, cplx: bool| {
1366 stats::enable(true);
1367 let _ = stats::take();
1368 if cplx {
1369 let a = Mat::<c64>::from_fn(m, k, |i, j| c64::new((i + j) as f64, 1.0));
1370 let b = Mat::<c64>::from_fn(k, n, |i, j| c64::new(i as f64, j as f64));
1371 let mut d = Mat::<c64>::zeros(m, n);
1372 matmul(d.as_mut(), Accum::Replace, a.as_ref(), b.as_ref(), c64::new(1.0, 0.0), par.reborrow());
1373 } else {
1374 let a = rand_mat(m, k, 1);
1375 let b = rand_mat(k, n, 2);
1376 let mut d = Mat::<f64>::zeros(m, n);
1377 matmul(d.as_mut(), Accum::Replace, a.as_ref(), b.as_ref(), 1.0, par.reborrow());
1378 }
1379 let st = stats::take();
1380 stats::enable(false);
1381 st
1382 };
1383 with_par(6, |mut par| {
1384 let st = before(64, 64, 64, &mut par, true);
1386 assert_eq!(st.gemm_below_floor, 0, "{st:?}");
1387 assert_eq!(engaged(&st), 1, "{st:?}");
1388 let st = before(20, 20, 20, &mut par, false);
1390 assert_eq!(engaged(&st), 0, "{st:?}");
1391 let st = before(192, 192, 192, &mut par, false);
1393 assert_eq!(engaged(&st), 1, "{st:?}");
1394 });
1395 }
1396
1397 #[test]
1398 fn teams_of_odd_and_unit_width_match_sequential() {
1399 let (m, n, k) = (300, 200, 250);
1400 let a = rand_mat(m, k, 1);
1401 let b = rand_mat(k, n, 2);
1402 let mut want = Mat::<f64>::zeros(m, n);
1403 matmul(want.as_mut(), Accum::Replace, a.as_ref(), b.as_ref(), 1.0, Par::Seq);
1404 for width in [1usize, 3, 5] {
1405 let mut got = Mat::<f64>::zeros(m, n);
1406 with_par(width, |par| {
1407 matmul(got.as_mut(), Accum::Replace, a.as_ref(), b.as_ref(), 1.0, par)
1408 });
1409 assert!(max_diff(want.as_ref(), got.as_ref()) < 1e-12, "width = {width}");
1410 }
1411 }
1412
1413 #[test]
1414 fn for_each_deals_every_item_once_and_nests() {
1415 with_par(6, |mut par| {
1416 for jobs in [1usize, 2, 3, 6, 9] {
1417 let hits: Vec<AtomicUsize> = (0..13).map(|_| AtomicUsize::new(0)).collect();
1418 crate::utils::thread::for_each_with_par(
1419 par.reborrow(),
1420 jobs,
1421 rayon::iter::IntoParallelIterator::into_par_iter(0..13usize),
1422 |i, child| {
1423 hits[i].fetch_add(1, Ordering::Relaxed);
1424 assert!(child.degree() <= 6);
1426 },
1427 );
1428 assert!(hits.iter().all(|h| h.load(Ordering::Relaxed) == 1), "jobs = {jobs}");
1429 }
1430 });
1431 }
1432
1433 #[test]
1434 fn join_raw_leaves_match_the_width() {
1435 for n in [2usize, 3, 5, 6] {
1436 let widths = std::sync::Mutex::new(Vec::new());
1437 with_par(n, |par| {
1438 crate::utils::thread::join_raw(
1439 |p| widths.lock().unwrap().push(p.degree()),
1440 |p| widths.lock().unwrap().push(p.degree()),
1441 par,
1442 )
1443 });
1444 let mut w = widths.into_inner().unwrap();
1445 w.sort_unstable();
1446 let half = n / 2;
1447 assert_eq!(w, vec![half, n - half], "n = {n}");
1448 }
1449 }
1450
1451 #[test]
1452 fn lu_qr_cholesky_on_the_team_match_sequential() {
1453 use crate::linalg::solvers::{Llt, Qr};
1454 use crate::Side;
1455 with_par(6, |mut par| {
1456 for n in [200usize, 513] {
1457 let a = rand_mat(n, n, 7);
1458 let b = rand_mat(n, 5, 8);
1459 let lu = |par: Par| {
1461 let mut f = a.clone();
1462 let mut perm = vec![0usize; n];
1463 let mut perm_inv = vec![0usize; n];
1464 let mut mem = dyn_stack::MemBuffer::new(
1465 crate::linalg::lu::partial_pivoting::factor::lu_in_place_scratch::<usize, f64>(n, n, par.shape(), Default::default()),
1466 );
1467 crate::linalg::lu::partial_pivoting::factor::lu_in_place(
1468 f.as_mut(),
1469 &mut perm,
1470 &mut perm_inv,
1471 par,
1472 dyn_stack::MemStack::new(&mut mem),
1473 Default::default(),
1474 );
1475 f
1476 };
1477 let want = lu(Par::Seq);
1478 let got = lu(par.reborrow());
1479 assert!(max_diff(want.as_ref(), got.as_ref()) < 1e-8, "lu n = {n}");
1480 let mut r_want = a.clone();
1482 let mut r_got = a.clone();
1483 let qr = |m: &mut Mat<f64>, par: Par| {
1484 let bs = crate::linalg::qr::no_pivoting::factor::recommended_block_size::<f64>(n, n);
1485 let mut h = Mat::<f64>::zeros(bs, n);
1486 let mut mem = dyn_stack::MemBuffer::new(
1487 crate::linalg::qr::no_pivoting::factor::qr_in_place_scratch::<f64>(n, n, bs, par.shape(), Default::default()),
1488 );
1489 crate::linalg::qr::no_pivoting::factor::qr_in_place(
1490 m.as_mut(),
1491 h.as_mut(),
1492 par,
1493 dyn_stack::MemStack::new(&mut mem),
1494 Default::default(),
1495 );
1496 };
1497 qr(&mut r_want, Par::Seq);
1498 qr(&mut r_got, par.reborrow());
1499 assert!(max_diff(r_want.as_ref(), r_got.as_ref()) < 1e-8, "qr n = {n}");
1500 let _ = (&b, Side::Lower);
1501 let _: Option<(Llt<f64>, Qr<f64>)> = None;
1502 let spd = {
1504 let mut s = Mat::<f64>::zeros(n, n);
1505 matmul(s.as_mut(), Accum::Replace, a.transpose(), a.as_ref(), 1.0, Par::Seq);
1506 for i in 0..n {
1507 s[(i, i)] += n as f64;
1508 }
1509 s
1510 };
1511 let ch = |par: Par| {
1512 let mut f = spd.clone();
1513 let mut mem = dyn_stack::MemBuffer::new(
1514 crate::linalg::cholesky::llt::factor::cholesky_in_place_scratch::<f64>(n, par.shape(), Default::default()),
1515 );
1516 crate::linalg::cholesky::llt::factor::cholesky_in_place(
1517 f.as_mut(),
1518 Default::default(),
1519 par,
1520 dyn_stack::MemStack::new(&mut mem),
1521 Default::default(),
1522 )
1523 .unwrap();
1524 f
1525 };
1526 let want = ch(Par::Seq);
1527 let got = ch(par.reborrow());
1528 let mut d: f64 = 0.0;
1530 for j in 0..n {
1531 for i in j..n {
1532 d = d.max((want[(i, j)] - got[(i, j)]).abs());
1533 }
1534 }
1535 assert!(d < 1e-8, "cholesky n = {n}: {d}");
1536 }
1537 });
1538 }
1539
1540 fn counted<R>(width: usize, f: impl FnOnce(Par<'_>) -> R) -> (R, stats::Stats) {
1542 with_par(width, |par| {
1543 stats::enable(true);
1544 let _ = stats::take();
1545 let r = f(par);
1546 let st = stats::take();
1547 stats::enable(false);
1548 (r, st)
1549 })
1550 }
1551
1552 #[test]
1553 fn sparse_matmul_on_a_team_runs_on_the_team() {
1554 use crate::sparse::{SparseColMat, Triplet};
1555 let n = 3000;
1556 let triplets: Vec<_> = (0..n)
1557 .flat_map(|j| {
1558 (0..4).map(move |t| {
1559 let i = (j * 7 + t * 131) % n;
1560 Triplet::new(i, j, ((i + 2 * j) as f64).sin())
1561 })
1562 })
1563 .collect();
1564 let a = SparseColMat::<usize, f64>::try_new_from_triplets(n, n, &triplets).unwrap();
1565 let want = crate::sparse::linalg::matmul::sparse_sparse_matmul(
1566 a.as_ref(), a.as_ref(), 1.0, Par::Seq,
1567 )
1568 .unwrap();
1569 let (got, st) = counted(4, |par| {
1570 crate::sparse::linalg::matmul::sparse_sparse_matmul(a.as_ref(), a.as_ref(), 1.0, par).unwrap()
1571 });
1572 assert!(st.split.get(4).copied().unwrap_or(0) >= 1, "{:?}", st.split);
1574 assert!(max_diff(want.to_dense().as_ref(), got.to_dense().as_ref()) < 1e-12);
1575 }
1576
1577 #[test]
1578 fn lu_swap_stage_splits_the_team_it_is_handed() {
1579 let n = 513;
1580 let a = rand_mat(n, n, 11);
1581 let lu = |par: Par| {
1582 let mut f = a.clone();
1583 let mut perm = vec![0usize; n];
1584 let mut perm_inv = vec![0usize; n];
1585 let mut mem = dyn_stack::MemBuffer::new(
1586 crate::linalg::lu::partial_pivoting::factor::lu_in_place_scratch::<usize, f64>(
1587 n, n, par.shape(), Default::default(),
1588 ),
1589 );
1590 crate::linalg::lu::partial_pivoting::factor::lu_in_place(
1591 f.as_mut(), &mut perm, &mut perm_inv, par,
1592 dyn_stack::MemStack::new(&mut mem), Default::default(),
1593 );
1594 f
1595 };
1596 let want = lu(Par::Seq);
1597 let (got, st) = counted(6, lu);
1598 assert!(st.split.get(2).copied().unwrap_or(0) >= 1, "{:?}", st.split);
1600 assert!(max_diff(want.as_ref(), got.as_ref()) < 1e-8);
1601 }
1602
1603 #[test]
1604 fn join_sized_on_a_team_hands_out_exactly_the_team() {
1605 use rayon::iter::IntoParallelIterator;
1606 for left in 0..=7usize {
1607 let active = AtomicUsize::new(0);
1608 let peak = AtomicUsize::new(0);
1609 let widths = [AtomicUsize::new(0), AtomicUsize::new(0)];
1610 with_par(6, |par| {
1611 let side = |which: usize, par: Par<'_>| {
1612 let jobs = par.degree();
1613 widths[which].store(jobs, Ordering::Relaxed);
1614 crate::utils::thread::for_each_with_par(par, jobs, (0..jobs).into_par_iter(), |_, _| {
1615 let now = active.fetch_add(1, Ordering::AcqRel) + 1;
1616 peak.fetch_max(now, Ordering::AcqRel);
1617 std::thread::sleep(std::time::Duration::from_millis(3));
1618 active.fetch_sub(1, Ordering::AcqRel);
1619 });
1620 };
1621 crate::utils::thread::join_sized(par, left, |p| side(0, p), |p| side(1, p));
1622 });
1623 assert_eq!(
1624 widths[0].load(Ordering::Relaxed) + widths[1].load(Ordering::Relaxed),
1625 6,
1626 "left = {left}"
1627 );
1628 assert!(peak.load(Ordering::Relaxed) <= 6, "left = {left}");
1629 }
1630 }
1631 #[test]
1632 fn a_team_of_one_behaves_exactly_like_seq() {
1633 use crate::utils::cost::{par_width, rate, ParThreshold};
1634 let _turn = team_turn();
1636 threadteam::Exec::serial()
1637 .team(|team| {
1638 let par = Par::Team(team);
1639 assert_eq!(par.degree(), 1);
1640 assert_eq!(par_width::<f64>(&par, 1 << 40, rate::gemm_mac(), 1), 1);
1642 assert!(!crate::utils::cost::join_pays::<f64>(&par, 1 << 40, rate::gemm_mac(), 1));
1643 assert_eq!(ParThreshold::Derived.region_width::<f64>(&par, 1 << 40, rate::gemm_mac(), 1), 1);
1644 let a = rand_mat(200, 150, 1);
1646 let b = rand_mat(150, 180, 2);
1647 let mut want = Mat::<f64>::zeros(200, 180);
1648 let mut got = Mat::<f64>::zeros(200, 180);
1649 matmul(want.as_mut(), Accum::Replace, a.as_ref(), b.as_ref(), 1.0, Par::Seq);
1650 stats::enable(true);
1651 let _ = stats::take();
1652 matmul(got.as_mut(), Accum::Replace, a.as_ref(), b.as_ref(), 1.0, par);
1653 let st = stats::take();
1654 stats::enable(false);
1655 assert_eq!(max_diff(want.as_ref(), got.as_ref()), 0.0);
1656 assert_eq!(st.gemm_team_kernel + st.gemm_tiled, 0, "{st:?}");
1657 })
1658 .unwrap();
1659 }
1660
1661 #[test]
1665 fn triangular_products_on_the_team_match_sequential_and_stay_on_it() {
1666 use crate::linalg::matmul::triangular::{matmul as tri_matmul, BlockStructure};
1667 let (n, k) = (768, 768);
1668 let a = rand_mat(n, k, 5);
1669 let b = rand_mat(k, n, 6);
1670 let lower = |par: Par| {
1671 let mut c = rand_mat(n, n, 7);
1672 tri_matmul(c.as_mut(), BlockStructure::TriangularLower, Accum::Add, a.as_ref(), BlockStructure::Rectangular, b.as_ref(), BlockStructure::Rectangular, 0.5, par);
1673 c
1674 };
1675 let upper = |par: Par| {
1676 let mut c = rand_mat(n, n, 8);
1677 tri_matmul(c.as_mut(), BlockStructure::TriangularUpper, Accum::Replace, a.as_ref(), BlockStructure::Rectangular, b.as_ref(), BlockStructure::Rectangular, 1.0, par);
1678 c
1679 };
1680 let (want_l, want_u) = (lower(Par::Seq), upper(Par::Seq));
1681 let ((got_l, got_u), st) = counted(6, |mut par| (lower(par.reborrow()), upper(par.reborrow())));
1682 let tri = |m: &Mat<f64>, w: &Mat<f64>, lo: bool| {
1683 let mut d: f64 = 0.0;
1684 for j in 0..n {
1685 for i in 0..n {
1686 if (lo && i >= j) || (!lo && i <= j) {
1687 d = d.max((m[(i, j)] - w[(i, j)]).abs());
1688 }
1689 }
1690 }
1691 d
1692 };
1693 assert!(tri(&got_l, &want_l, true) < 1e-9);
1694 assert!(tri(&got_u, &want_u, false) < 1e-9);
1695 assert!(st.gemm_tiled > 0 || st.gemm_team_kernel > 0 || st.split.iter().skip(2).any(|&c| c > 0), "{st:?}");
1697 }
1698
1699 #[test]
1702 fn spicy_matmul_on_a_team_matches_sequential() {
1703 use crate::linalg::matmul::internal::spicy_matmul;
1704 use crate::linalg::matmul::triangular::BlockStructure;
1705 if !crate::linalg::matmul::internal::has_spicy_matmul::<f64>() {
1706 return;
1707 }
1708 let (n, k) = (768, 512);
1709 let a = rand_mat(n, k, 9);
1710 let b = rand_mat(k, n, 10);
1711 let run = |par: Par| {
1712 let mut c = rand_mat(n, n, 11);
1713 spicy_matmul::<usize, f64>(c.as_mut(), BlockStructure::TriangularLower, None, None, Accum::Add, a.as_ref(), crate::Conj::No, b.as_ref(), crate::Conj::No, None, 0.5, par, dyn_stack::MemStack::new(&mut []));
1714 c
1715 };
1716 let want = run(Par::Seq);
1717 let got = with_par(6, |par| run(par));
1718 let mut d: f64 = 0.0;
1719 for j in 0..n {
1720 for i in j..n {
1721 d = d.max((want[(i, j)] - got[(i, j)]).abs());
1722 }
1723 }
1724 assert!(d < 1e-9, "{d}");
1725 }
1726
1727 #[test]
1730 fn gemm_tiled_matches_the_gemm_crate_for_every_dtype() {
1731 use num_complex::Complex;
1732 #[cfg(all(target_arch = "x86_64", feature = "std"))]
1733 assert_eq!(
1734 x86_kernel::<f64>().is_some(),
1735 std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("fma")
1736 );
1737 fn check<E: 'static + Copy + Send + Sync + core::fmt::Debug>(
1738 par: &mut Par<'_>,
1739 one: E,
1740 from: impl Fn(f64) -> E,
1741 err: impl Fn(E, E) -> f64,
1742 cplx: bool,
1743 ) {
1744 let (m, n, k) = (301, 263, 190);
1745 let gen = |len: usize, seed: u64| -> Vec<E> {
1746 let mut x = seed.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
1747 (0..len)
1748 .map(|_| {
1749 x = x.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
1750 from(((x >> 11) as f64 / (1u64 << 53) as f64) - 0.5)
1751 })
1752 .collect()
1753 };
1754 let (a, b) = (gen(m * k, 1), gen(k * n, 2));
1755 let zero = from(0.0);
1756 let alpha = from(1.25);
1757 for read_dst in [false, true] {
1758 for conj in [false, true] {
1759 if conj && !cplx {
1760 continue;
1761 }
1762 let base = gen(m * n, 3);
1763 let (mut want, mut got) = (base.clone(), base);
1764 let acc = if read_dst { one } else { zero };
1765 unsafe {
1766 gemm::gemm(
1767 m, n, k, want.as_mut_ptr(), m as isize, 1, read_dst,
1768 a.as_ptr(), m as isize, 1, b.as_ptr(), k as isize, 1,
1769 acc, alpha, false, conj, conj, gemm::Parallelism::None,
1770 );
1771 let Par::Team(team) = par.reborrow() else { unreachable!() };
1772 gemm_tiled::<E>(
1773 team, m, n, k, 1e3, got.as_mut_ptr(), m as isize, 1, read_dst,
1774 a.as_ptr(), m as isize, 1, b.as_ptr(), k as isize, 1,
1775 acc, alpha, conj, conj,
1776 );
1777 }
1778 let d = want.iter().zip(&got).map(|(&w, &g)| err(w, g)).fold(0.0, f64::max);
1779 assert!(d < 1e-3, "read_dst {read_dst} conj {conj}: {d}");
1780 }
1781 }
1782 }
1783 with_par(5, |mut par| {
1784 check::<f64>(&mut par, 1.0, |x| x, |w, g| (w - g).abs(), false);
1785 check::<f32>(&mut par, 1.0, |x| x as f32, |w, g| (w - g).abs() as f64, false);
1786 check::<Complex<f64>>(&mut par, Complex::new(1.0, 0.0), |x| Complex::new(x, 0.5 - x), |w, g| (w - g).norm(), true);
1787 check::<Complex<f32>>(&mut par, Complex::new(1.0, 0.0), |x| Complex::new(x as f32, 0.5 - x as f32), |w, g| (w - g).norm() as f64, true);
1788 });
1789 let (_, st) = counted(5, |mut par| check::<f64>(&mut par, 1.0, |x| x, |w, g| (w - g).abs(), false));
1790 assert!(st.gemm_tiled >= 2 && st.gemm_below_floor == 0, "{st:?}");
1791 }
1792
1793 #[cfg(all(target_arch = "x86_64", feature = "std"))]
1795 struct Recording<'a> {
1796 inner: &'a dyn private_gemm_x86::Fork,
1797 seen: std::sync::Mutex<Vec<std::thread::ThreadId>>,
1798 runs: AtomicUsize,
1799 }
1800
1801 #[cfg(all(target_arch = "x86_64", feature = "std"))]
1802 impl private_gemm_x86::Fork for Recording<'_> {
1803 fn width(&self) -> usize {
1804 self.inner.width()
1805 }
1806 fn run(&self, f: &(dyn Fn(usize) + Sync)) {
1807 self.runs.fetch_add(1, Ordering::Relaxed);
1808 let seen = &self.seen;
1809 self.inner.run(&|j| {
1810 seen.lock().unwrap().push(std::thread::current().id());
1811 f(j)
1812 });
1813 }
1814 }
1815
1816 #[cfg(all(target_arch = "x86_64", feature = "std"))]
1821 #[test]
1822 fn the_kernel_forks_only_on_members_of_the_team() {
1823 use private_gemm_x86::{gemm_with, Accum as KAccum, DType, DstKind, IType, InstrSet};
1824 let Some((feat, _)) = super::kernel::available::<f64>() else { return };
1825 let _ = InstrSet::Avx256;
1826 let (m, n, k) = (400usize, 360usize, 300usize);
1827 let a = rand_mat(m, k, 1);
1828 let b = rand_mat(k, n, 2);
1829 let mut want = Mat::<f64>::zeros(m, n);
1830 matmul(want.as_mut(), Accum::Replace, a.as_ref(), b.as_ref(), 1.0, Par::Seq);
1831 let _turn = team_turn();
1832 let team = team_of(6);
1833 let mut exec = team.exec(6);
1834 exec.team(|mut token| {
1835 let members = std::sync::Mutex::new(Vec::new());
1837 token.reborrow().run(|_, _| members.lock().unwrap().push(std::thread::current().id()));
1838 let members = members.into_inner().unwrap();
1839 assert_eq!(members.len(), 6);
1840 for width in [6usize, 3, 2] {
1841 let mut got = Mat::<f64>::zeros(m, n);
1842 let alpha = 1.0f64;
1843 let rec = std::cell::RefCell::new(None);
1844 super::kernel::with_fork(token.reborrow(), width, |fork| {
1845 assert_eq!(fork.width(), width);
1846 let recording = Recording { inner: fork, seen: Default::default(), runs: AtomicUsize::new(0) };
1847 unsafe {
1848 gemm_with(
1849 DType::F64, IType::U64, feat, m, n, k,
1850 got.as_mut().as_ptr_mut() as *mut (), 1, got.col_stride(),
1851 core::ptr::null(), core::ptr::null(), DstKind::Full, KAccum::Replace,
1852 a.as_ref().as_ptr() as *const (), 1, a.col_stride(), false,
1853 core::ptr::null(), 0,
1854 b.as_ref().as_ptr() as *const (), 1, b.col_stride(), false,
1855 &raw const alpha as *const (),
1856 &recording,
1857 );
1858 }
1859 *rec.borrow_mut() = Some((recording.seen.into_inner().unwrap(), recording.runs.into_inner()));
1860 });
1861 let (seen, runs) = rec.into_inner().unwrap();
1862 assert!(runs >= 1, "width {width}: the kernel never forked");
1863 assert_eq!(seen.len(), runs * width, "width {width}");
1864 assert!(seen.iter().all(|t| members.contains(t)), "width {width}: a participant outside the team");
1865 let distinct: std::collections::HashSet<_> = seen.iter().collect();
1866 assert_eq!(distinct.len(), width, "width {width}: seen {} entries over {runs} runs", seen.len());
1867 assert!(max_diff(want.as_ref(), got.as_ref()) < 1e-10, "width {width}");
1868 }
1869 })
1870 .ok()
1871 .expect("a team token");
1872 }
1873
1874 #[cfg(all(target_arch = "x86_64", feature = "std"))]
1877 #[test]
1878 fn team_kernel_products_match_sequential_for_every_dtype() {
1879 use num_complex::Complex;
1880 if super::kernel::available::<f64>().is_none() {
1881 return;
1882 }
1883 fn check<T: crate::ComplexField + Copy>(from: impl Fn(f64, f64) -> T, err: impl Fn(T, T) -> f64) {
1884 let (m, n, k) = (333usize, 287usize, 211usize);
1885 let gen = |r: usize, c: usize, seed: u64| {
1886 let mut x = seed.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
1887 Mat::<T>::from_fn(r, c, |_, _| {
1888 x = x.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
1889 let u = (x >> 11) as f64 / (1u64 << 53) as f64 - 0.5;
1890 from(u, 0.25 - u)
1891 })
1892 };
1893 let (a, b) = (gen(m, k, 1), gen(k, n, 2));
1894 let alpha = from(0.75, -0.25);
1895 for beta in [Accum::Replace, Accum::Add] {
1896 let base = gen(m, n, 3);
1897 let (mut want, mut got) = (base.clone(), base);
1898 matmul(want.as_mut(), beta, a.as_ref(), b.as_ref(), alpha, Par::Seq);
1899 let ((), st) = counted(6, |par| matmul(got.as_mut(), beta, a.as_ref(), b.as_ref(), alpha, par));
1900 let mut d = 0.0f64;
1901 for j in 0..n {
1902 for i in 0..m {
1903 d = d.max(err(want[(i, j)], got[(i, j)]));
1904 }
1905 }
1906 assert!(d < 1e-3, "{beta:?}: {d}");
1907 assert_eq!(st.gemm_team_kernel, 1, "{st:?}");
1908 assert_eq!(st.gemm_tiled, 0, "{st:?}");
1909 }
1910 }
1911 check::<f64>(|x, _| x, |w, g| (w - g).abs());
1912 check::<f32>(|x, _| x as f32, |w, g| (w - g).abs() as f64);
1913 check::<Complex<f64>>(|x, y| Complex::new(x, y), |w, g| (w - g).norm());
1914 check::<Complex<f32>>(|x, y| Complex::new(x as f32, y as f32), |w, g| (w - g).norm() as f64);
1915 }
1916
1917 #[cfg(all(target_arch = "x86_64", feature = "std"))]
1921 #[test]
1922 fn team_kernel_triangular_indexed_and_factorizations_match_sequential() {
1923 use crate::linalg::matmul::internal::spicy_matmul;
1924 use crate::linalg::matmul::triangular::{matmul as tri_matmul, BlockStructure};
1925 if !crate::linalg::matmul::internal::has_spicy_matmul::<f64>() {
1926 return;
1927 }
1928 let (n, k) = (700usize, 500usize);
1929 let a = rand_mat(n, k, 21);
1930 let b = rand_mat(k, n, 22);
1931 let tri_diff = |x: &Mat<f64>, y: &Mat<f64>, lower: bool| {
1932 let mut d: f64 = 0.0;
1933 for j in 0..n {
1934 for i in 0..n {
1935 if (lower && i >= j) || (!lower && i <= j) {
1936 d = d.max((x[(i, j)] - y[(i, j)]).abs());
1937 }
1938 }
1939 }
1940 d
1941 };
1942 for (block, lower) in [(BlockStructure::TriangularLower, true), (BlockStructure::TriangularUpper, false)] {
1944 let run = |par: Par| {
1945 let mut c = rand_mat(n, n, 23);
1946 spicy_matmul::<usize, f64>(c.as_mut(), block, None, None, Accum::Add, a.as_ref(), crate::Conj::No, b.as_ref(), crate::Conj::No, None, 0.5, par, dyn_stack::MemStack::new(&mut []));
1947 c
1948 };
1949 let want = run(Par::Seq);
1950 let (got, st) = counted(6, run);
1951 assert!(tri_diff(&want, &got, lower) < 1e-9);
1952 assert_eq!(st.gemm_team_kernel, 1, "{st:?}");
1953 }
1954 {
1956 let rows: Vec<usize> = (0..n).map(|i| (i * 3) % (3 * n - 1) + (i % 2)).collect::<std::collections::BTreeSet<_>>().into_iter().take(n).collect();
1957 let cols: Vec<usize> = (0..n).map(|j| 2 * j).collect();
1958 let (nr, nc) = (rows.iter().max().unwrap() + 1, 2 * n);
1959 let m_rows = rows.len();
1960 let a2 = rand_mat(m_rows, k, 24);
1961 let run = |par: Par| {
1962 let mut c = rand_mat(nr, nc, 25);
1963 spicy_matmul::<usize, f64>(c.as_mut(), BlockStructure::Rectangular, Some(&rows), Some(&cols), Accum::Add, a2.as_ref(), crate::Conj::No, b.as_ref(), crate::Conj::No, None, 0.5, par, dyn_stack::MemStack::new(&mut []));
1964 c
1965 };
1966 let want = run(Par::Seq);
1967 let (got, st) = counted(6, run);
1968 assert!(max_diff(want.as_ref(), got.as_ref()) < 1e-9);
1969 assert_eq!(st.gemm_team_kernel, 1, "{st:?}");
1970 }
1971 {
1973 let run = |par: Par| {
1974 let mut c = rand_mat(n, n, 26);
1975 tri_matmul(c.as_mut(), BlockStructure::TriangularLower, Accum::Add, a.as_ref(), BlockStructure::Rectangular, b.as_ref(), BlockStructure::Rectangular, 0.5, par);
1976 c
1977 };
1978 let want = run(Par::Seq);
1979 let (got, st) = counted(6, run);
1980 assert!(tri_diff(&want, &got, true) < 1e-9);
1981 assert!(st.gemm_team_kernel >= 1, "{st:?}");
1982 }
1983 let n = 640usize;
1985 let g = rand_mat(n, n, 30);
1986 let spd = {
1987 let mut s = Mat::<f64>::zeros(n, n);
1988 matmul(s.as_mut(), Accum::Replace, g.transpose(), g.as_ref(), 1.0, Par::Seq);
1989 for i in 0..n {
1990 s[(i, i)] += n as f64;
1991 }
1992 s
1993 };
1994 let lower_diff = |x: &Mat<f64>, y: &Mat<f64>| {
1995 let mut d: f64 = 0.0;
1996 for j in 0..n {
1997 for i in j..n {
1998 d = d.max((x[(i, j)] - y[(i, j)]).abs());
1999 }
2000 }
2001 d
2002 };
2003 let llt = |par: Par| {
2004 let mut f = spd.clone();
2005 let mut mem = dyn_stack::MemBuffer::new(crate::linalg::cholesky::llt::factor::cholesky_in_place_scratch::<f64>(n, par.shape(), Default::default()));
2006 crate::linalg::cholesky::llt::factor::cholesky_in_place(f.as_mut(), Default::default(), par, dyn_stack::MemStack::new(&mut mem), Default::default()).unwrap();
2007 f
2008 };
2009 let (got, st) = counted(6, llt);
2010 assert!(lower_diff(&llt(Par::Seq), &got) < 1e-8, "llt");
2011 assert!(st.gemm_team_kernel > 0, "llt: {st:?}");
2012 let ldlt = |par: Par| {
2013 let mut f = spd.clone();
2014 let mut mem = dyn_stack::MemBuffer::new(crate::linalg::cholesky::ldlt::factor::cholesky_in_place_scratch::<f64>(n, par.shape(), Default::default()));
2015 crate::linalg::cholesky::ldlt::factor::cholesky_in_place(
2016 f.as_mut(),
2017 crate::linalg::cholesky::ldlt::factor::LdltRegularization { dynamic_regularization_signs: None, dynamic_regularization_delta: 0.0, dynamic_regularization_epsilon: 0.0 },
2018 par,
2019 dyn_stack::MemStack::new(&mut mem),
2020 Default::default(),
2021 )
2022 .unwrap();
2023 f
2024 };
2025 let (got, st) = counted(6, ldlt);
2026 assert!(lower_diff(&ldlt(Par::Seq), &got) < 1e-8, "ldlt");
2027 assert!(st.gemm_team_kernel > 0, "ldlt: {st:?}");
2028 let qr = |par: Par| {
2029 let mut f = g.clone();
2030 let bs = crate::linalg::qr::no_pivoting::factor::recommended_block_size::<f64>(n, n);
2031 let mut h = Mat::<f64>::zeros(bs, n);
2032 let mut mem = dyn_stack::MemBuffer::new(crate::linalg::qr::no_pivoting::factor::qr_in_place_scratch::<f64>(n, n, bs, par.shape(), Default::default()));
2033 crate::linalg::qr::no_pivoting::factor::qr_in_place(f.as_mut(), h.as_mut(), par, dyn_stack::MemStack::new(&mut mem), Default::default());
2034 f
2035 };
2036 let got = with_par(6, qr);
2039 assert!(max_diff(qr(Par::Seq).as_ref(), got.as_ref()) < 1e-8, "qr");
2040 let lu = |par: Par| {
2041 let mut f = g.clone();
2042 let (mut perm, mut perm_inv) = (vec![0usize; n], vec![0usize; n]);
2043 let mut mem = dyn_stack::MemBuffer::new(crate::linalg::lu::partial_pivoting::factor::lu_in_place_scratch::<usize, f64>(n, n, par.shape(), Default::default()));
2044 crate::linalg::lu::partial_pivoting::factor::lu_in_place(f.as_mut(), &mut perm, &mut perm_inv, par, dyn_stack::MemStack::new(&mut mem), Default::default());
2045 f
2046 };
2047 let (got, st) = counted(6, lu);
2048 assert!(max_diff(lu(Par::Seq).as_ref(), got.as_ref()) < 1e-8, "lu");
2049 assert!(st.gemm_team_kernel > 0, "lu: {st:?}");
2050 }
2051
2052 #[cfg(all(target_arch = "x86_64", feature = "std"))]
2054 #[test]
2055 fn team_kernel_qr_matches_sequential() {
2056 if super::kernel::available::<f64>().is_none() {
2057 return;
2058 }
2059 let n = 1400usize;
2060 let g = rand_mat(n, n, 31);
2061 let qr = |par: Par| {
2062 let mut f = g.clone();
2063 let bs = crate::linalg::qr::no_pivoting::factor::recommended_block_size::<f64>(n, n);
2064 let mut h = Mat::<f64>::zeros(bs, n);
2065 let mut mem = dyn_stack::MemBuffer::new(crate::linalg::qr::no_pivoting::factor::qr_in_place_scratch::<f64>(n, n, bs, par.shape(), Default::default()));
2066 crate::linalg::qr::no_pivoting::factor::qr_in_place(f.as_mut(), h.as_mut(), par, dyn_stack::MemStack::new(&mut mem), Default::default());
2067 f
2068 };
2069 let (got, st) = counted(6, qr);
2070 assert!(max_diff(qr(Par::Seq).as_ref(), got.as_ref()) < 1e-7, "qr");
2071 assert!(st.gemm_team_kernel > 0, "qr: {st:?}");
2072 }
2073
2074 #[test]
2077 fn high_level_singular_values_run_on_the_team() {
2078 let n = 384;
2079 let a = rand_mat(n, n, 21);
2080 let want = a.as_ref().singular_values(Par::Seq).unwrap();
2081 let (got, st) = counted(4, |par| a.as_ref().singular_values(par).unwrap());
2082 assert_eq!(want.len(), got.len());
2083 for (w, g) in want.iter().zip(&got) {
2084 assert!((w - g).abs() < 1e-9 * want[0], "{w} {g}");
2085 }
2086 let opened = st.run.iter().skip(2).sum::<usize>() + st.split.iter().skip(2).sum::<usize>();
2087 assert!(opened > 0 || st.gemm_tiled > 0 || st.gemm_team_kernel > 0, "no team section: {st:?}");
2088 }
2089
2090 #[test]
2092 fn high_level_qr_and_lu_solve_on_the_team() {
2093 use crate::linalg::solvers::Solve;
2094 let n = 320;
2095 let a = rand_mat(n, n, 22);
2096 let b = rand_mat(n, 8, 23);
2097 let want = a.as_ref().partial_piv_lu(Par::Seq).solve(b.as_ref(), Par::Seq);
2098 let (got, st) = counted(4, |mut par| {
2099 let lu = a.as_ref().partial_piv_lu(par.reborrow());
2100 lu.solve(b.as_ref(), par)
2101 });
2102 assert!(max_diff(want.as_ref(), got.as_ref()) < 1e-8);
2103 let opened = st.run.iter().skip(2).sum::<usize>() + st.split.iter().skip(2).sum::<usize>();
2104 assert!(opened > 0 || st.gemm_tiled > 0 || st.gemm_team_kernel > 0, "no team section: {st:?}");
2105 let want = a.as_ref().qr(Par::Seq).solve(b.as_ref(), Par::Seq);
2106 let got = with_par(4, |mut par| {
2107 let qr = a.as_ref().qr(par.reborrow());
2108 qr.solve(b.as_ref(), par)
2109 });
2110 assert!(max_diff(want.as_ref(), got.as_ref()) < 1e-8);
2111 }
2112
2113 #[test]
2117 fn with_team_calls_of_one_width_reuse_the_threads() {
2118 use std::collections::HashSet;
2119 use std::sync::Mutex;
2120 let _turn = team_turn();
2121 let me = std::thread::current().id();
2122 let workers = || {
2124 with_team(5, |par| {
2125 let Par::Team(mut team) = par else { panic!("a width of 5 is a team") };
2126 let ids = Mutex::new(HashSet::new());
2127 team.reborrow().run(|_, _| {
2128 ids.lock().unwrap().insert(std::thread::current().id());
2129 });
2130 let mut ids = ids.into_inner().unwrap();
2131 ids.remove(&me);
2132 ids
2133 })
2134 };
2135 let first = workers();
2136 let second = workers();
2137 assert_eq!((first.len(), second.len()), (4, 4));
2138 let all = first.union(&second).count();
2139 assert!(all <= 5, "two calls ran on {all} threads: the second pool spawned new ones");
2140 }
2141}