#![allow(missing_docs)]
pub mod bound;
pub mod cost;
#[cfg(feature = "threadteam")]
pub mod team;
#[doc(hidden)]
pub mod thread {
use crate::*;
#[inline]
pub fn join_raw(
op_a: impl Send + FnOnce(Par),
op_b: impl Send + FnOnce(Par),
parallelism: Par,
) {
fn implementation<'a>(
op_a: &'a mut (dyn Send + FnMut(Par)),
op_b: &'a mut (dyn Send + FnMut(Par)),
mut parallelism: Par,
) {
match parallelism {
Par::Seq => {
(op_a(parallelism.reborrow()), op_b(parallelism));
},
#[cfg(feature = "rayon")]
Par::Rayon(n_threads) => {
let n_threads = n_threads.get();
if n_threads == 1 {
(op_a(Par::Seq), op_b(Par::Seq));
} else {
let par = |n: usize| {
Par::Rayon(core::num::NonZeroUsize::new(n).unwrap())
};
let half = n_threads / 2;
spindle::for_each(
2,
[(op_a, par(n_threads - half)), (op_b, par(half))],
|(op, parallelism)| op(parallelism),
);
}
},
#[cfg(feature = "threadteam")]
Par::Team(team) => {
let n_threads = team.size().get();
if n_threads == 1 {
(op_a(Par::Team(team)), op_b(Par::Seq));
} else {
let half = n_threads / 2;
let ops = [
std::sync::Mutex::new(op_a),
std::sync::Mutex::new(op_b),
];
super::team::note_split(2);
let _ = team.split(&[n_threads - half, half], |c, child| {
(ops[c].lock().unwrap())(Par::Team(child))
});
}
},
};
}
let mut op_a = Some(op_a);
let mut op_b = Some(op_b);
implementation(
&mut |parallelism| op_a.take().unwrap()(parallelism),
&mut |parallelism| op_b.take().unwrap()(parallelism),
parallelism,
)
}
pub struct Ptr<T>(pub *mut T);
unsafe impl<T> Send for Ptr<T> {}
unsafe impl<T> Sync for Ptr<T> {}
impl<T> Copy for Ptr<T> {}
impl<T> Clone for Ptr<T> {
#[inline]
fn clone(&self) -> Self {
*self
}
}
#[inline]
pub fn parallelism_degree(parallelism: &Par) -> usize {
match parallelism {
Par::Seq => 1,
#[cfg(feature = "rayon")]
Par::Rayon(n_threads) => n_threads.get(),
#[cfg(feature = "threadteam")]
Par::Team(team) => team.size().get(),
#[cfg(not(feature = "threadteam"))]
Par::__Lifetime(never, _) => match *never {},
}
}
#[inline]
pub fn is_team(parallelism: &Par) -> bool {
match parallelism {
#[cfg(feature = "threadteam")]
Par::Team(_) => true,
_ => false,
}
}
#[cfg(feature = "rayon")]
pub fn par_for_each<
T: Send,
I: rayon::iter::IntoParallelIterator<Item = T>,
>(
par: Par,
iter: I,
f: impl Send + Sync + Fn(T),
) where
I::Iter: rayon::iter::IndexedParallelIterator,
{
match par {
#[cfg(feature = "threadteam")]
Par::Team(team) => {
let width = team.size().get();
super::team::for_each(team, width, iter, f)
},
_ => {
use rayon::iter::ParallelIterator;
iter.into_par_iter().for_each(f)
},
}
}
#[cfg(feature = "rayon")]
pub fn for_each<
T: Send,
I: rayon::iter::IntoParallelIterator<Item = T>,
>(
par: Par,
n_jobs: usize,
iter: I,
f: impl Sync + Fn(T),
) where
I::Iter: rayon::iter::IndexedParallelIterator,
{
match par {
Par::Seq => {
use rayon::iter::plumbing::{Producer, ProducerCallback};
struct Each<F>(F);
impl<T, F: Fn(T)> ProducerCallback<T> for Each<F> {
type Output = ();
fn callback<P: Producer<Item = T>>(self, p: P) {
p.into_iter().for_each(self.0);
}
}
rayon::iter::IndexedParallelIterator::with_producer(
iter.into_par_iter(),
Each(f),
);
},
Par::Rayon(_) => spindle::for_each(n_jobs, iter, f),
#[cfg(feature = "threadteam")]
Par::Team(team) => super::team::for_each(team, n_jobs, iter, f),
}
}
#[cfg(feature = "rayon")]
pub fn for_each_with_par<
T: Send,
I: rayon::iter::IntoParallelIterator<Item = T>,
>(
par: Par,
n_jobs: usize,
iter: I,
f: impl Sync + Fn(T, Par),
) where
I::Iter: rayon::iter::IndexedParallelIterator,
{
match par {
Par::Seq => for_each(Par::Seq, n_jobs, iter, |item| f(item, Par::Seq)),
Par::Rayon(n) => {
let width = n.get() / n_jobs.max(1);
spindle::for_each(n_jobs, iter, |item| {
f(item, if width >= 2 { Par::rayon(width) } else { Par::Seq })
})
},
#[cfg(feature = "threadteam")]
Par::Team(team) => {
super::team::for_each_child(team, n_jobs, iter, |item, child| {
f(item, Par::Team(child))
})
},
}
}
#[cfg(feature = "rayon")]
pub fn join_sized(
par: Par,
size_a: usize,
op_a: impl Send + FnOnce(Par),
op_b: impl Send + FnOnce(Par),
) {
let degree = par.degree();
if degree < 2 {
op_a(Par::Seq);
op_b(Par::Seq);
return;
}
let size_a = size_a.clamp(1, degree - 1);
let op_a = std::sync::Mutex::new(Some(op_a));
let op_b = std::sync::Mutex::new(Some(op_b));
split(par, &[size_a, degree - size_a], &|c, child| {
if c == 0 {
(op_a.lock().unwrap().take().unwrap())(child)
} else {
(op_b.lock().unwrap().take().unwrap())(child)
}
});
}
#[cfg(feature = "rayon")]
pub fn split(par: Par, sizes: &[usize], f: &(dyn Sync + Fn(usize, Par))) {
core::assert!(sizes.iter().all(|&s| s > 0));
match par {
Par::Seq => {
for c in 0..sizes.len() {
f(c, Par::Seq);
}
},
Par::Rayon(n) => {
core::assert!(sizes.iter().sum::<usize>() <= n.get());
use rayon::iter::IntoParallelIterator;
spindle::for_each(sizes.len(), (0..sizes.len()).into_par_iter(), |c| {
f(c, if sizes[c] >= 2 { Par::rayon(sizes[c]) } else { Par::Seq })
})
},
#[cfg(feature = "threadteam")]
Par::Team(team) => {
if sizes.iter().sum::<usize>() <= team.size().get() {
super::team::note_split(sizes.len());
let _ = team.split(sizes, |c, child| f(c, Par::Team(child)));
} else {
for c in 0..sizes.len() {
f(c, Par::Seq);
}
}
},
}
}
#[inline]
pub fn par_split_indices(
n: usize,
idx: usize,
chunk_count: usize,
) -> (usize, usize) {
let chunk_size = n / chunk_count;
let rem = n % chunk_count;
let idx_to_col_start = move |idx| {
if idx < rem {
idx * (chunk_size + 1)
} else {
rem + idx * chunk_size
}
};
let start = idx_to_col_start(idx);
let end = idx_to_col_start(idx + 1);
(start, end - start)
}
#[cfg(all(test, feature = "rayon"))]
mod tests {
use super::{for_each_with_par, join_raw, join_sized, parallelism_degree};
use crate::Par;
use core::sync::atomic::{AtomicUsize, Ordering};
struct Leaves {
count: AtomicUsize,
active: AtomicUsize,
peak: AtomicUsize,
}
fn split(parallelism: Par, leaves: &Leaves) {
if parallelism_degree(¶llelism) > 1 {
join_raw(
|par| split(par, leaves),
|par| split(par, leaves),
parallelism,
);
} else {
leaves.count.fetch_add(1, Ordering::Relaxed);
let now = leaves.active.fetch_add(1, Ordering::AcqRel) + 1;
leaves.peak.fetch_max(now, Ordering::AcqRel);
std::thread::sleep(std::time::Duration::from_millis(5));
leaves.active.fetch_sub(1, Ordering::AcqRel);
}
}
#[test]
fn join_raw_leaves_do_not_exceed_requested_width() {
let pool = rayon::ThreadPoolBuilder::new()
.num_threads(16)
.build()
.unwrap();
for n in [2usize, 3, 5, 6, 7, 8] {
let leaves = Leaves {
count: AtomicUsize::new(0),
active: AtomicUsize::new(0),
peak: AtomicUsize::new(0),
};
pool.install(|| {
split(
Par::Rayon(core::num::NonZeroUsize::new(n).unwrap()),
&leaves,
)
});
assert_eq!(leaves.count.into_inner(), n, "leaves for n = {n}");
assert!(
leaves.peak.into_inner() <= n,
"concurrent leaves exceed n = {n}"
);
}
}
#[test]
fn join_sized_children_stay_within_their_share() {
use rayon::prelude::*;
let pool = rayon::ThreadPoolBuilder::new()
.num_threads(16)
.build()
.unwrap();
for n in [2usize, 3, 5, 8] {
for left in 0..=n + 1 {
let active = AtomicUsize::new(0);
let peak = AtomicUsize::new(0);
let widths = [AtomicUsize::new(0), AtomicUsize::new(0)];
let side = |which: usize, par: Par| {
widths[which].store(parallelism_degree(&par), Ordering::Relaxed);
let jobs = parallelism_degree(&par);
for_each_with_par(par, jobs, (0..jobs).into_par_iter(), |_, _| {
let now = active.fetch_add(1, Ordering::AcqRel) + 1;
peak.fetch_max(now, Ordering::AcqRel);
std::thread::sleep(std::time::Duration::from_millis(3));
active.fetch_sub(1, Ordering::AcqRel);
});
};
pool.install(|| {
join_sized(
Par::Rayon(core::num::NonZeroUsize::new(n).unwrap()),
left,
|par| side(0, par),
|par| side(1, par),
)
});
let total = widths[0].load(Ordering::Relaxed)
+ widths[1].load(Ordering::Relaxed);
assert_eq!(total, n, "n = {n}, left = {left}");
assert!(
peak.load(Ordering::Relaxed) <= n,
"concurrent jobs exceed n = {n} with left = {left}"
);
}
}
}
}
}
pub mod approx;
pub mod simd;