rustymath_core/reduce/

sliding.rs

use super::*;

/// The widest window that [`sliding`] folds directly, with one pass over each block of a lane for each of the window's rows. Wider windows take van Herk-Gil-Werman.
pub const DIRECT: usize = 8;
/// The items in each block of a lane that [`sliding`]'s direct folds finish before moving on, so that a block and the rows its windows read stay in cache.
const SLIDE_BLOCK: usize = 1024;

/// Van Herk-Gil-Werman's folds of the windows of `width` rows in `lane`, rows of `inner` items, into `folds`. `suffix` holds the lane's block suffixes, and `row` carries one row's running fold. `ONE` fixes `inner` at 1, which compiles each loop over a row's items to scalar code.
#[inline(always)]
fn hgw<T: Copy, R: Reduction<T>, const ONE: bool>(
    lane: &[T], folds: &mut [R::Total], suffix: &mut [R::Total], row: &mut [R::Total], width: usize, inner: usize, ok: &mut bool) {
    let inner = if ONE { 1 } else { inner };
    // A local row of one item stays in a register, where a row in memory would put a store and a load in each step's chain. The loops index the row item by item: `steps_of` put it in memory and made `⌈/50↕v` 1.4 times slower.
    let mut one = [R::unit()];
    let row = if ONE { &mut one[..] } else { &mut row[..inner] };
    for (b, block) in lane.chunks(width * inner).enumerate() {
        let (start, last) = (b * width * inner, block.len() - inner);
        // Each row's suffix folds it with the rest of its block, from the block's last row back to its first. Each step of a suffix or a prefix waits for the one before, so both combine through `chain`, the earlier items on the left.
        let s = &mut suffix[start..start + block.len()];
        for j in 0..inner {
            row[j] = one_of::<T, R>(block[last + j], ok);
            s[last + j] = row[j];
        }
        for k in (0..last).step_by(inner).rev() {
            for j in 0..inner {
                row[j] = chain_of::<T, R>(one_of::<T, R>(block[k + j], ok), row[j], ok);
                s[k + j] = row[j];
            }
        }
        // Each row's prefix folds the block up to it. The first block's whole prefix is the first window.
        for j in 0..inner { row[j] = one_of::<T, R>(block[j], ok); }
        if b == 0 {
            for k in (inner..block.len()).step_by(inner) { for j in 0..inner { row[j] = chain_of::<T, R>(row[j], one_of::<T, R>(block[k + j], ok), ok); } }
            folds[..inner].copy_from_slice(row);
            continue;
        }
        // In a later block, row `i` ends the window that starts `i + 1` rows into the block before: the fold of that block's suffix there and this prefix.
        let from = start + inner - width * inner;
        let (windows, earlier) = (&mut folds[from..from + block.len()], &suffix[from..from + block.len()]);
        for j in 0..inner { windows[j] = merge_of::<T, R>(earlier[j], row[j], ok); }
        for k in (inner..block.len()).step_by(inner) {
            for j in 0..inner {
                row[j] = chain_of::<T, R>(row[j], one_of::<T, R>(block[k + j], ok), ok);
                windows[k + j] = merge_of::<T, R>(earlier[k + j], row[j], ok);
            }
        }
        // A window that starts this block is the whole block, which the prefix holds once its last row is in.
        if block.len() == width * inner { folds[start..start + inner].copy_from_slice(row); }
    }
}

/// The folds by `R` of each window of `width` neighbouring rows along the middle axis of `data`, an array of shape `[outer, len, inner]` held row by row. The result has shape `[outer, len - width + 1, inner]`, with whether every fold fits. Each fold combines its window's items in order, except that a window as long as a lane of single items folds as [`along`] folds the lane. A window of at most [`DIRECT`] rows folds directly. A wider one takes van Herk-Gil-Werman's method: each lane splits into blocks of `width` rows, and each window is the fold of one block's suffix and the next block's prefix, so an item costs the same whatever the width.
pub fn sliding<T: Copy + Sync, R: Reduction<T>>(data: &[T], [outer, len, inner]: [usize; 3], width: usize) -> (Vec<R::Total>, bool) {
    assert_eq!(data.len(), outer * len * inner, "the data must fill its dimensions");
    assert!((1..=len.max(1)).contains(&width), "a window lies within its lane");
    let finished = |(totals, fits)| {
        if fits { crate::events::taken(crate::events::Route::SlidingFold, data.len()); }
        (totals, fits)
    };
    let count = (len + 1).saturating_sub(width);
    if outer * count * inner == 0 { return finished((Vec::new(), true)); }
    // A window as long as a lane of single items is that lane's fold, which `along` runs side by side with its neighbours'. Wider lanes fold directly, starting from their first row.
    if width == len && inner == 1 { return finished(along::<T, R>(data, [outer, len, inner])); }
    let mut ok = true;
    if width <= DIRECT {
        // Each block of a lane's windows folds while its rows are in cache (`rows_folded`), which reads the lane from memory once.
        let block = (SLIDE_BLOCK / inner).max(1);
        let mut out = Vec::with_capacity(outer * count * inner);
        for lane in data.chunks_exact(len * inner) {
            for r in (0..count).step_by(block) {
                let rows = block.min(count - r);
                ok &= rows_folded::<T, R>(&mut out, width, |i| &lane[(r + i) * inner..(r + i + rows) * inner]);
            }
        }
        return finished((out, ok));
    }
    let mut out = vec![R::unit(); outer * count * inner];
    let lanes = data.chunks_exact(len * inner).zip(out.chunks_exact_mut(count * inner));
    let (mut suffix, mut row) = (vec![R::unit(); len * inner], vec![R::unit(); inner]);
    for (lane, folds) in lanes {
        if inner == 1 { hgw::<T, R, true>(lane, folds, &mut suffix, &mut row, width, inner, &mut ok) }
        else { hgw::<T, R, false>(lane, folds, &mut suffix, &mut row, width, inner, &mut ok) }
    }
    finished((out, ok))
}

/// The folds by `R` of each window of `w1` rows along axis 1 and `w2` rows along axis 3 of `data`, an array of shape `[outer, len1, mid, len2, inner]` held row by row. Windows start every `e1` positions along axis 1 and every `e2` along axis 3: every position for windows that slide, or every window's width for windows that tile the axis, as pooling's do. The result has shape `[outer, c1, mid, c2, inner]`, where `c1` and `c2` count the windows along each axis, with whether every fold fits. Each band of `w1` rows folds into a buffer that stays in cache, and the windows along axis 3 fold from that buffer, so `data` is read once. A window's items fold along axis 1 and then along axis 3, which suits reductions that fold in any order. `None` for a window wider than [`DIRECT`] rows.
pub fn sliding_2d<T: Copy + Sync, R: Reduction<T>>(
    data: &[T], [outer, len1, mid, len2, inner]: [usize; 5], [w1, w2]: [usize; 2], [e1, e2]: [usize; 2]) -> Option<(Vec<R::Total>, bool)> {
    assert_eq!(data.len(), outer * len1 * mid * len2 * inner, "the data must fill its dimensions");
    assert!(w1 > 0 && w2 > 0 && e1 > 0 && e2 > 0, "windows hold items and start in order");
    if w1 > DIRECT || w2 > DIRECT { return None; }
    let (c1, c2) = ((len1 + e1).saturating_sub(w1) / e1, (len2 + e2).saturating_sub(w2) / e2);
    if outer * c1 * mid * c2 * inner == 0 { return Some((Vec::new(), true)); }
    let (row, plane) = (len2 * inner, mid * len2 * inner);
    let mut out = Vec::with_capacity(outer * c1 * mid * c2 * inner);
    let (mut band, mut ok) = (Vec::with_capacity(plane), true);
    for lane in data.chunks_exact(len1 * plane) {
        for i in 0..c1 {
            band.clear();
            ok &= rows_folded::<T, R>(&mut band, w1, |a| &lane[(e1 * i + a) * plane..(e1 * i + a + 1) * plane]);
            for m in band.chunks_exact(row) {
                let at = out.len();
                // Windows that slide fold as `sliding`'s direct folds do, one pass over the band's row for each of the window's rows. Windows that tile fold one window at a time.
                if e2 == 1 {
                    out.extend_from_slice(&m[..c2 * inner]);
                    for b in 1..w2 { ok &= merges_of::<T, R>(&mut out[at..], &m[b * inner..]); }
                } else {
                    for j in 0..c2 { out.extend_from_slice(&m[e2 * j * inner..(e2 * j + 1) * inner]); }
                    for b in 1..w2 {
                        for j in 0..c2 { ok &= merges_of::<T, R>(&mut out[at + j * inner..at + (j + 1) * inner], &m[(e2 * j + b) * inner..]); }
                    }
                }
            }
        }
    }
    if ok { crate::events::taken(crate::events::Route::SlidingFold, data.len()); }
    Some((out, ok))
}

/// A reduction that picks one of its items, as [`Max`] and [`Min`] do, over pairs of an item and its row. Each fold keeps the row of the first item that gives it.
struct Pick<R>(std::marker::PhantomData<R>);

impl<T: Copy + PartialEq + Send + Sync, R: Reduction<T, Total = T>> Reduction<(T, usize)> for Pick<R> {
    type Total = (T, usize);
    #[inline(always)]
    fn unit() -> (T, usize) { (R::unit(), 0) }
    #[inline(always)]
    fn one((x, row): (T, usize)) -> ((T, usize), bool) { ((R::one(x).0, row), true) }
    #[inline(always)]
    fn step(acc: (T, usize), x: (T, usize)) -> (T, usize) { Self::merge(acc, Self::one(x).0).0 }
    /// The earlier fold `a` where it gives the fold of both, and `b` where `b`'s item is beyond it.
    #[inline(always)]
    fn merge(a: (T, usize), b: (T, usize)) -> ((T, usize), bool) { (if R::merge(a.0, b.0).0 == a.0 { a } else { b }, true) }
    fn slice(x: &[(T, usize)]) -> ((T, usize), bool) { (x.iter().fold(Self::unit(), |t, &v| Self::step(t, v)), true) }
    fn scan(x: &[(T, usize)], start: Option<(T, usize)>, out: &mut Vec<(T, usize)>) -> bool {
        let mut total = start;
        for &v in x {
            let t = total.map_or_else(|| Self::one(v).0, |t| Self::step(t, v));
            out.push(t);
            total = Some(t);
        }
        true
    }
}

/// The folds that [`picks`] finds the extremes of: windows of a width that slide along each lane, as [`sliding`] folds them, each whole lane, as [`along`] folds it, or each lane's running folds, as [`scan_along`] gives them.
#[derive(Clone, Copy, Debug)]
pub enum Picks { Windows(usize), Lanes, Running }

/// The row of the first item that gives each fold by `R` along the middle axis of `data`, an array of shape `[outer, len, inner]`, for a reduction that picks one of its items, as [`Max`] and [`Min`] do. `of` chooses the folds, and the result lays out as their results do. Each row counts from the start of its lane. A fold whose items all leave it at `R::unit()`, as NaN items do for a float maximum, gives its first row.
pub fn picks<T: Copy + PartialEq + Send + Sync, R: Reduction<T, Total = T>>(data: &[T], shape: [usize; 3], of: Picks) -> Vec<usize> {
    let [outer, len, inner] = shape;
    match of {
        Picks::Windows(width) => {
            let pairs: Vec<(T, usize)> = data.iter().enumerate().map(|(k, &x)| (x, k / inner.max(1) % len.max(1))).collect();
            sliding::<(T, usize), Pick<R>>(&pairs, shape, width).0.into_iter().map(|(_, row)| row).collect()
        }
        // A lane's pick is its first item that, folded alone, gives the lane's fold. Along a lane of consecutive items, a vectorized search finds it after the fold.
        Picks::Lanes if inner == 1 => along::<T, R>(data, shape).0.into_iter().enumerate().map(|(o, total)| {
            let row = first_where(&data[o * len..][..len], |x| R::one(x).0 == total);
            if row == len { 0 } else { row }
        }).collect(),
        // Lanes across an axis fold in one pass in memory order, and each lane's pick steps through `moved`.
        Picks::Lanes => {
            let mut rows = vec![0; outer * inner];
            if len > 0 && inner > 0 {
                for (block, lanes) in data.chunks_exact(len * inner).zip(rows.chunks_exact_mut(inner)) {
                    let mut totals: Vec<T> = block[..inner].iter().map(|&x| R::one(x).0).collect();
                    for (k, items) in block.chunks_exact(inner).enumerate().skip(1) {
                        for ((row, total), &x) in lanes.iter_mut().zip(&mut totals).zip(items) {
                            let next = R::merge(*total, R::one(x).0).0;
                            *row = moved(*row, *total, next, k);
                            *total = next;
                        }
                    }
                }
            }
            rows
        }
        // A running fold's pick at each item steps through `moved` from the item before it in the lane.
        Picks::Running => {
            let mut totals = Vec::with_capacity(data.len());
            scan_along::<T, R>(data, shape, None, &mut totals);
            let mut rows = vec![0; data.len()];
            for at in 0..data.len() { let k = at / inner % len; if k > 0 { rows[at] = moved(rows[at - inner], totals[at - inner], totals[at], k) } }
            rows
        }
    }
}

/// [`picks`] of the items of `x` at their own type, by [`Max`], or by [`Min`] where `max` is false. `None` for complex numbers and characters, which have no order.
pub fn view_picks(x: View<'_>, dims: [usize; 3], of: Picks, max: bool) -> Option<Vec<usize>> {
    fn picked<T: Copy + PartialEq + Send + Sync>(x: &[T], dims: [usize; 3], of: Picks, max: bool) -> Vec<usize> where Max: Reduction<T, Total = T>,
        Min: Reduction<T, Total = T> { if max { picks::<T, Max>(x, dims, of) } else { picks::<T, Min>(x, dims, of) } }
    Some(match x {
        View::Booleans(x) => picked(x, dims, of, max), View::Integers(x) => with_ints!(x, |x| picked(x, dims, of, max)),
        View::Floats(x) => with_floats!(x, |x| picked(x, dims, of, max)), _ => return None
    })
}