rustymath_core/

parallel.rs

//! 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) }
}