//! Splitting work across threads. [`pieces`] cuts a call's items into pieces of the sizes its caller gives, and [`ranges`] into as many pieces as the time of its first unit justifies. The pieces run on a Rayon pool with Rayon's default number of threads: `RAYON_NUM_THREADS` where it's set, and otherwise one for each core. [`set_threads`] changes that number for the whole process. A call made on one of a pool's threads, such as a function that another call's piece runs, doesn't split again. In the browser build, which has no threads, every call runs its pieces in turn on the calling thread.
use crate::slice::{Arg, CHUNK};
use crate::{error::Error, interrupt::check};
/// The clock that `parallel`'s rounds and the speed tests' timings read: the browser's in a browser build, where `std::time::Instant` panics.
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
pub use std::time::Instant;
use std::{mem::MaybeUninit, ops::Range, sync::atomic::{AtomicUsize, Ordering::Relaxed}, time::Duration};
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
pub use web_time::Instant;
/// The number of threads that calls split across, for the whole process, or 0 for the default pool's number.
static THREADS: AtomicUsize = AtomicUsize::new(0);
/// Sets how many threads calls split their work across, for the whole process, and gives the number it replaces. 0 restores the default. 1 runs every call on its calling thread, as the speed tests do for their budgets.
pub fn set_threads(n: usize) -> usize { THREADS.swap(n, Relaxed) }
/// The number that [`set_threads`] last set. It starts at 0, the default.
pub fn thread_setting() -> usize { THREADS.load(Relaxed) }
/// How many threads the calls made on this thread split their work across: the process's [`set_threads`] count, or 1 on a pool's worker thread.
pub fn threads() -> usize { requested().map_or(1, |n| if n == 0 { pool::get(0).current_num_threads() } else { n }) }
/// Whether calls made on this thread run serially: after `set_threads(1)`, or on a pool's worker thread.
pub(crate) fn serial() -> bool { requested().is_none() }
/// The pool that calls made on this thread split across, as its number of threads, 0 for the default. `None` where they run on this thread.
fn requested() -> Option<usize> { let n = THREADS.load(Relaxed); (n != 1 && !pool::inside()).then_some(n) }
/// Runs `piece` on each range that cuts `0..n` into pieces of `size` items, the last one shorter, and gives whether every call said `true`. Every piece runs, whatever another one gives. The pieces run across at most `most` of this thread's [`threads`], on a smaller pool where `most` or the number of pieces is below the pool's threads, as [`sized`] chooses. Where either is 1, they run in order on this thread.
pub(crate) fn pieces(n: usize, size: usize, most: usize, piece: &(dyn Fn(Range<usize>) -> bool + Sync)) -> bool {
let (count, range) = (n.div_ceil(size), |p: usize| p * size..n.min((p + 1) * size));
if count.min(most) > 1 && let Some(threads) = requested() { return pool::run(sized(threads, count.min(most)), count, &|p| piece(range(p))); }
(0..count).fold(true, |ok, p| piece(range(p)) & ok)
}
/// The pool that `count` pieces run on, where this thread's calls split across the pool of `threads` threads, 0 for the default. Where there are fewer pieces than that pool's threads, it's the pool of the largest power of two up to `count`. On an M5 Max, `-v` on a million items, cut into about 4 pieces, took 0.33 copies on a 4-thread pool and 0.5 to 0.8 on pools of 12 to 18 threads.
fn sized(threads: usize, count: usize) -> usize {
let all = if threads == 0 { pool::get(0).current_num_threads() } else { threads };
if count >= all { threads } else { 1 << count.ilog2() }
}
/// Runs `write` on each piece of `size` items that cuts `out`, the last one shorter, across at most `most` threads as [`pieces`] runs its ranges, and gives whether every call said `true`. `write(start, items)` writes `items`, the piece of `out` from `start`.
pub fn pieces_of<B: Send>(out: &mut [B], size: usize, most: usize, write: &(dyn Fn(usize, &mut [B]) -> bool + Sync)) -> bool {
let (shared, n) = (Shared(out.as_mut_ptr()), out.len());
// SAFETY: `pieces` gives each call its own range within `0..n`.
pieces(n, size, most, &|range| write(range.start, unsafe { shared.part(range) }))
}
/// An output slice shared only through disjoint ranges.
pub struct Shared<B>(*mut B);
// SAFETY: each piece reaches only the items in its own range, so no two threads reach the same item.
unsafe impl<B: Send> Sync for Shared<B> {}
impl<B> Shared<B> {
/// The items in `range`. SAFETY: the range must lie within the output, and no other caller may hold any of these items at the same time.
#[inline]
unsafe fn part(&self, range: Range<usize>) -> &mut [B] { unsafe { std::slice::from_raw_parts_mut(self.0.add(range.start), range.len()) } }
}
/// The time on one thread that each piece of a [`ranges`] call is given at least. A call whose pieces would be shorter splits into fewer of them, because waking a thread for a shorter piece costs more than the piece saves. On an M5 Max, `v+w` on a million items, about 200 µs on one thread, took 0.51 copies split across 4 threads and 0.81 across all 18.
pub const PIECE_TIME: Duration = Duration::from_micros(30);
/// Runs `run` on disjoint ranges covering `0..n`, and gives whether every call said `true`. Every range runs, even after a failure. Each range contains whole units of `unit` items, except that the last unit may be shorter. A call of at most one unit, or one whose calls run on this thread, runs whole on this thread. A longer call times its first unit on this thread, then cuts the rest into as many pieces as that estimate gives [`PIECE_TIME`] each, at most four for each of this thread's [`threads`]. It runs them as [`pieces`] does, or runs the rest in one call on this thread where that gives fewer than two pieces. This non-generic, out-of-line scheduler keeps Rayon's machinery out of its callers. An empty call runs `0..0` once. Panics if `unit` is zero.
#[inline(never)]
pub fn ranges(n: usize, unit: usize, run: &(dyn Fn(Range<usize>) -> bool + Sync)) -> bool {
assert!(unit > 0, "a work unit must contain at least one item");
if n <= unit || serial() { return run(0..n); }
let start = Instant::now();
let ok = run(0..unit);
let units = (n - unit).div_ceil(unit);
let count = ((start.elapsed().as_secs_f64() * units as f64 / PIECE_TIME.as_secs_f64()) as usize).min(4 * threads()).min(units);
if count < 2 { return ok & run(unit..n); }
ok & pieces(n - unit, units.div_ceil(count) * unit, usize::MAX, &|range| run(unit + range.start..unit + range.end))
}
/// Runs `write(start, items)` over disjoint pieces of `out`, using [`ranges`] in units of `unit` items, and gives whether every call said `true`. An empty output calls `write` once. Panics if `unit` is zero.
#[inline(never)]
pub fn split<B: Send>(out: &mut [B], unit: usize, write: &(dyn Fn(usize, &mut [B]) -> bool + Sync)) -> bool {
let (shared, n) = (Shared(out.as_mut_ptr()), out.len());
// SAFETY: `ranges` gives each call its own range within `0..n`.
ranges(n, unit, &|range| write(range.start, unsafe { shared.part(range) }))
}
/// Adds the sums of every piece of `0..n` into `sums`, where `piece(range, sums)` adds the sums of the piece `range` into the `sums` it's given. `0..n` splits into one piece of at least `least` items for each of this thread's [`threads`]. With too few items for two pieces, `0..n` runs as one piece on this thread. Otherwise the first piece adds into `sums`, each later one into its own part of `parts`, and the parts then add into `sums` on this thread.
pub fn summed<T: Copy + Send + num_traits::Zero>(n: usize, least: usize, sums: &mut [T], parts: &mut Vec<T>, piece: &(dyn Fn(Range<usize>, &mut [T]) + Sync)) {
let pieces = (n / least).min(threads());
if pieces <= 1 { return piece(0..n, sums); }
let per = n.div_ceil(pieces);
let (len, width) = (sums.len(), sums.len().max(1));
parts.clear();
parts.resize((n.div_ceil(per) - 1) * width, T::zero());
{
let mut work: Vec<&mut [T]> = std::iter::once(&mut *sums).chain(parts.chunks_mut(width).map(|part| &mut part[..len])).collect();
pieces_of(&mut work, 1, usize::MAX, &|p, work| { piece(p * per..n.min((p + 1) * per), &mut work[0]); true });
}
for part in parts.chunks(width) { for (s, &x) in sums.iter_mut().zip(part) { *s = *s + x } }
}
/// [`summed`] over the units of `items`, each of `unit` items, where `piece(first, items, sums)` gets the items of the piece whose first unit is `first`.
pub fn summed_of<B: Send, T: Copy + Send + num_traits::Zero>(
items: &mut [B], unit: usize, least: usize, sums: &mut [T], parts: &mut Vec<T>, piece: &(dyn Fn(usize, &mut [B], &mut [T]) + Sync)) {
let shared = Shared(items.as_mut_ptr());
// SAFETY: `summed` gives each call its own range of units, within `items`.
summed(items.len() / unit, least, sums, parts, &|units, sums| piece(units.start, unsafe { shared.part(units.start * unit..units.end * unit) }, sums))
}
/// The most cells that [`each_cell`] runs alone at its start, to time a cell.
const CELLS: usize = 8;
/// The pieces that each thread is given in a round of an [`each_cell`] call, so that a thread that finishes early takes more of the round.
const PIECES: usize = 4;
/// The time that each round of an [`each_cell`] call gives each thread, where its cells are that short. A stack stops within about a round of an interrupt, or within one cell where a cell takes longer.
const ROUND_TIME: Duration = Duration::from_millis(20);
/// Runs `write(i, scratch)` on each cell `i` of a stack of `count` cells, and gives whether every call said `true`. Cells that run in turn on one thread share one `S` for their scratch buffers, such as an `LuWork`. It checks for an interrupt on this thread between rounds, and an interrupt's error stops the call. It makes no check before the first cell, which a caller's own check covers, so a stack of one small cell costs no check.
///
/// The first cells run alone on this thread, to time a cell: up to `CELLS` of them, or [`PIECE_TIME`] of work, because one cell alone reads high from a cold start. Each later round holds about `ROUND_TIME` of cells for each of this thread's [`threads`], at least one each. A round's cells split into about `PIECES` pieces for each thread, each of at least `PIECE_TIME` of work. A cell that runs on this thread can split its own work, and one on another thread doesn't. On all 18 threads of an M5 Max, splitting stacks of 8 to 16 matrices of 128² to 512² was 1.5 to 5 times as fast as running their cells in turn, each split across the threads. Three inverses of 1024² were 25% slower. In the browser build, which has one thread, each round's cells run in turn. A stack of one cell runs alone, with nothing to time.
fn each_cell<S: Default>(count: usize, write: impl Fn(usize, &mut S) -> bool + Sync) -> Result<bool, Error> {
let mut scratch = S::default();
if count <= 1 { return Ok(count == 0 || write(0, &mut scratch)); }
let threads = threads();
let start = Instant::now();
let mut ok = write(0, &mut scratch);
// The timed cells after a short first cell skip the clock read between them. A stack of 64 tiny products spent about a sixth of its time reading the clock.
let short = start.elapsed() * CELLS as u32 <= PIECE_TIME;
let mut timed = 1;
while timed < count.min(CELLS) && (short || start.elapsed() < PIECE_TIME) {
ok &= write(timed, &mut scratch);
timed += 1;
}
let cell = start.elapsed().as_secs_f64().max(1e-9) / timed as f64;
let round = ((ROUND_TIME.as_secs_f64() / cell) as usize).max(1).saturating_mul(threads);
let least = (PIECE_TIME.as_secs_f64() / cell).ceil() as usize;
for first in (timed..count).step_by(round) {
check()?;
let cells = round.min(count - first);
if threads == 1 {
ok &= (first..first + cells).fold(true, |ok, i| write(i, &mut scratch) & ok);
} else {
ok &= pieces(cells, cells.div_ceil(threads * PIECES).max(least).max(1), usize::MAX, &|range| {
let mut scratch = S::default();
range.map(|c| first + c).fold(true, |ok, i| write(i, &mut scratch) & ok)
});
}
}
Ok(ok)
}
/// A stack's outputs: a tuple of one to three `Vec`s, each holding its own number of items for every cell in turn. [`stack`] gives each cell its part of each output, as a QR stack writes Q and R and an SVD stack U, S and V.
pub trait Outputs: Sized {
/// The items that each cell writes in each output.
type Per: Copy + Sync;
/// One cell's part of each output.
type Cell<'a>
where
Self: 'a;
/// The outputs' items, shared with the threads that write their cells.
type Shared: Sync;
/// The outputs of `count` cells, each item at its default. `vec!` asks the allocator for zeroed memory where that default is zero bits, as for real numbers and integers, so no pass writes a large output before its cells do.
fn new(count: usize, per: Self::Per) -> Self;
/// Pointers to the outputs' items.
fn shared(&mut self) -> Self::Shared;
/// Cell `i`'s part of each output. SAFETY: `i` must be below the outputs' number of cells, and no other caller may hold cell `i`'s part at the same time.
unsafe fn cell<'a>(shared: &'a Self::Shared, per: Self::Per, i: usize) -> Self::Cell<'a> where Self: 'a;
}
macro_rules! outputs {
($n:literal: $($t:ident $k:tt),+) => {
impl<$($t: Clone + Default + Send + 'static),+> Outputs for ($(Vec<$t>,)+) {
type Per = [usize; $n];
type Cell<'a> = ($(&'a mut [$t],)+) where Self: 'a;
type Shared = ($(Shared<$t>,)+);
fn new(count: usize, per: [usize; $n]) -> Self { ($(vec![<$t>::default(); count * per[$k]],)+) }
fn shared(&mut self) -> Self::Shared { ($(Shared(self.$k.as_mut_ptr()),)+) }
unsafe fn cell<'a>(shared: &'a Self::Shared, per: [usize; $n], i: usize) -> Self::Cell<'a> where Self: 'a {
// SAFETY: as the caller promises.
($(unsafe { shared.$k.part(i * per[$k]..(i + 1) * per[$k]) },)+)
}
}
};
}
outputs!(1: A 0);
outputs!(2: A 0, B 1);
outputs!(3: A 0, B 1, C 2);
/// Runs `write(i, parts, scratch)` on each cell `i` of a stack of `count` cells, as [`each_cell`] runs them, and gives the outputs with whether every call said `true`. `parts` is cell `i`'s part of each output `k`: `per[k]` items from `i * per[k]`.
pub fn stack<O: Outputs + 'static, S: Default>(
count: usize, per: O::Per, write: &(dyn Fn(usize, O::Cell<'_>, &mut S) -> bool + Sync)) -> Result<(O, bool), Error> {
let mut out = O::new(count, per);
let shared = out.shared();
// SAFETY: `each_cell` gives each cell below `count` to one call.
let ok = each_cell(count, |i, scratch| write(i, unsafe { O::cell(&shared, per, i) }, scratch))?;
Ok((out, ok))
}
/// `f(i, scratch)` of each cell `i` of a stack of `count` cells, in order, as [`each_cell`] runs them: [`stack`] for a result of any type, one for each cell.
pub fn each<R: Send, S: Default>(count: usize, f: impl Fn(usize, &mut S) -> R + Sync) -> Result<Vec<R>, Error> {
let mut out: Vec<Option<R>> = (0..count).map(|_| None).collect();
let shared = Shared(out.as_mut_ptr());
// SAFETY: `each_cell` gives each cell below `count` to one call.
each_cell(count, |i, scratch| {
let cell = unsafe { shared.part(i..i + 1) };
cell[0] = Some(f(i, scratch));
true
})?;
Ok(out.into_iter().map(|c| c.expect("`each_cell` runs every cell")).collect())
}
/// Appends `n` items to `out`, and gives whether every call of `write` said `true`. `write(range, slots)` fills `slots` with the items in `range`, in order. The calls split across threads as [`split`] splits a call, in units of `unit` items. A call that leaves any of its slots unfilled panics, and `out` keeps its length.
pub fn extended<B: Copy + Send>(out: &mut Vec<B>, n: usize, unit: usize, write: &(dyn Fn(Range<usize>, &mut Slots<'_, B>) -> bool + Sync)) -> bool {
unsafe {
crate::slice::appended(out, n, |slots| {
split(slots, unit, &|start, slots| {
let end = start + slots.len();
filled(slots, |slots| write(start..end, slots))
})
})
}
}
/// `write`'s result for `slots`, which it fills in order through one [`Slots`]. Panics where it leaves any unfilled.
fn filled<B, T>(slots: &mut [MaybeUninit<B>], write: impl FnOnce(&mut Slots<'_, B>) -> T) -> T {
let mut slots = Slots { slots, written: 0 };
let result = write(&mut slots);
assert_eq!(slots.written, slots.slots.len(), "a writer fills every slot it gets");
result
}
/// The slots of a vector that a writer fills in order: one call's part of an [`extended`] vector, or the items that [`Slots::append`] adds.
pub struct Slots<'a, B> { slots: &'a mut [MaybeUninit<B>], written: usize }
impl<B> Slots<'_, B> {
/// Appends `n` items to `out` on this thread, and gives `write`'s result. `write` fills the slots in order. Where it leaves any unfilled, `append` panics and `out` keeps its length.
pub fn append<T>(out: &mut Vec<B>, n: usize, write: impl FnOnce(&mut Slots<'_, B>) -> T) -> T {
unsafe { crate::slice::appended(out, n, |slots| filled(slots, write)) }
}
/// Writes the `n` results of [`pairs`](crate::pairs) after the items already written, on this thread, and gives whether every result succeeded. Panics where they don't fit.
pub fn pairs<X: Copy, Y: Copy>(&mut self, x: Arg<'_, X>, y: Arg<'_, Y>, n: usize, f: impl Fn(X, Y) -> (B, bool)) -> bool {
unsafe { self.fill(n, |slots| crate::slice::pairs_into(x, y, n, slots, f)) }
}
/// Writes `items` after the items already written, on this thread. Panics where they don't fit.
pub fn copy(&mut self, items: &[B]) where B: Copy {
unsafe {
self.fill(items.len(), |slots| {
for (slot, &item) in slots.iter_mut().zip(items) { slot.write(item); }
true
})
};
}
/// `write`'s result for the next `n` slots, which it must fill. They then count as written. Panics where they don't fit.
///
/// # Safety
/// `write` must initialize every slot it receives before returning, whether it returns `true` or `false`.
pub unsafe fn fill(&mut self, n: usize, write: impl FnOnce(&mut [MaybeUninit<B>]) -> bool) -> bool {
let ok = write(&mut self.slots[self.written..self.written + n]);
self.written += n;
ok
}
}
/// Runs `write` on each `CHUNK` of `slots` in turn, the results for the items from `start`, and gives whether every call said `true`. A slice form that makes several passes over buffers of its own keeps them in cache this way.
#[inline]
pub(crate) fn in_chunks<B>(start: usize, slots: &mut [MaybeUninit<B>], write: impl Fn(usize, &mut [MaybeUninit<B>]) -> bool) -> bool {
let (n, mut ok) = (slots.len(), true);
for at in (0..n).step_by(CHUNK) { ok &= write(start + at, &mut slots[at..n.min(at + CHUNK)]) }
ok
}
#[cfg(not(target_arch = "wasm32"))]
mod pool {
use std::sync::{Arc, Mutex};
/// The pools built in this process, each with its number of threads, 0 for the default, and the id of the process that built them. A child that `fork` creates has none of its parent's threads, so it builds pools of its own.
static POOLS: Mutex<(u32, Vec<(usize, Arc<rayon::ThreadPool>)>)> = Mutex::new((0, Vec::new()));
/// The pool of `n` threads, 0 for Rayon's default number, built the first time a call needs it.
pub(super) fn get(n: usize) -> Arc<rayon::ThreadPool> {
let mut pools = POOLS.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
let id = std::process::id();
if pools.0 != id {
// Dropping a parent's pool would signal threads that this process doesn't have.
std::mem::forget(std::mem::take(&mut pools.1));
pools.0 = id;
}
if let Some((_, pool)) = pools.1.iter().find(|(k, _)| *k == n) { return pool.clone(); }
let pool = Arc::new(rayon::ThreadPoolBuilder::new().num_threads(n).build().expect("a thread pool"));
pools.1.push((n, pool.clone()));
pool
}
/// Whether this thread is one of a pool's.
pub(super) fn inside() -> bool { rayon::current_thread_index().is_some() }
/// Runs `f` on each of `count` pieces across the pool of `threads` threads, and gives whether every call said `true`.
pub(super) fn run(threads: usize, count: usize, f: &(dyn Fn(usize) -> bool + Sync)) -> bool {
use rayon::prelude::*;
let error = crate::error::ErrorSlot::current();
get(threads).install(|| {
(0..count).into_par_iter().map(|item| match &error { Some(slot) => slot.within(|| f(item)), None => f(item) }).reduce(|| true, |a, b| a & b)
})
}
}
#[cfg(target_arch = "wasm32")]
mod pool {
pub(super) struct Pool;
impl Pool { pub(super) fn current_num_threads(&self) -> usize { 1 } }
pub(super) fn get(_: usize) -> Pool { Pool }
pub(super) fn inside() -> bool { true }
pub(super) fn run(_: usize, count: usize, f: &dyn Fn(usize) -> bool) -> bool { (0..count).fold(true, |ok, p| f(p) & ok) }
}