rustymath_core/reduce/

scans.rs

use super::*;

/// Appends the running folds by `R` along the middle axis of `data`, an array of shape `[outer, len, inner]` held row by row, to `out`, in the same layout. Where `seed` is given, each lane's first fold combines `seed`'s item for its column, or its only item. Gives whether every fold fits. Lanes of consecutive items scan by [`Reduction::short_lanes`] where the reduction takes them, else by `scan`. Along any other axis, each row's folds step the row before's by `checked`.
pub fn scan_along<T: Copy, R: Reduction<T>>(data: &[T], dims: [usize; 3], seed: Option<&[R::Total]>, out: &mut Vec<R::Total>) -> bool {
    let start = |i| seed.map(|s| s[if s.len() == 1 { 0 } else { i }]);
    if let [_, len, 1] = dims && let Some(fits) = R::short_lanes(data, len, start(0), out) { return fits; }
    rows_scan(data, dims, out, |lane, o| R::scan(lane, start(0), o), |i, x| start(i).map_or_else(|| R::one(x), |s| R::checked(s, x)), R::checked)
}

/// Appends the ordered left scan of each lane of `[outer, len, inner]` in `data` to `out`, in the same layout: along the last axis by `lane`, one call for every lane, stepping from `seed`'s item; along any other axis by `block`, one call for each row, stepping the row before it, or `seed`'s row, by the row's items. Without a seed, a lane's first item or row is kept as it is. Gives the kind of the first step that failed.
///
/// # Safety
/// `block` must initialize every output slot before returning, including when it returns `Err`.
pub unsafe fn ordered_scan<T: Copy>(
    data: &[T], [_, len, inner]: [usize; 3], seed: Option<&[T]>, out: &mut Vec<T>, lane: &dyn Fn(Option<T>, &[T], usize, &mut Vec<T>) -> Option<Kind>,
    block: &dyn Fn(&[T], &[T], &mut [std::mem::MaybeUninit<T>]) -> Result<(), Kind>) -> Option<Kind> {
    if inner == 1 { return lane(seed.map(|s| s[0]), data, len, out); }
    unsafe {
        appended(out, data.len(), |slots| {
            let mut failure = None;
            for (items, outs) in data.chunks(len * inner).zip(slots.chunks_mut(len * inner)) {
                match seed {
                    Some(s) => failure = failure.or(block(s, &items[..inner], &mut outs[..inner]).err()),
                    None => { for (slot, &x) in outs.iter_mut().zip(&items[..inner]) { slot.write(x); } }
                }
                for r in 1..len {
                    let (done, next) = outs.split_at_mut(r * inner);
                    failure = failure.or(block(written(&mut done[(r - 1) * inner..]), &items[r * inner..(r + 1) * inner], &mut next[..inner]).err());
                }
            }
            failure
        })
    }
}

/// Appends the running results along each lane of `len` consecutive items in `data` to `out`, and gives whether every one fits: `first` gives each lane's first result, and `step` each later one from the one before. Lanes shorter than [`SHORT_LANE`] run [`LANES`] side by side (`rows_scan`).
pub fn lanes_scan<T: Copy>(data: &[T], len: usize, out: &mut Vec<T>, first: impl Fn(T) -> (T, bool) + Copy, step: impl Fn(T, T) -> (T, bool) + Copy) -> bool {
    rows_scan(data, [0, len, 1], out, |lane, out| {
        let (total, ok) = first(lane[0]);
        out.push(total);
        checked_scan(&lane[1..], total, out, step) & ok
    }, move |_, x| first(x), step)
}

/// Appends the running results along the middle axis of `data`, an array of shape `[outer, len, inner]` held row by row, to `out`, in the same layout, and gives whether every one fits. A lane of at least [`SHORT_LANE`] consecutive items runs through `lane`. Shorter lanes run [`LANES`] side by side, and lanes along any other axis run a row at a time: `first` gives each lane's first item or row, and `step` steps each from the one before. `first` and `step` go to `side_by_side` by value, because a build of many codegen units can leave core's `<&F as Fn>::call` for a reference out of line, at a call for each item.
#[inline(always)]
fn rows_scan<T: Copy, S: Copy>(
    data: &[T], [_, len, inner]: [usize; 3], out: &mut Vec<S>, lane: impl Fn(&[T], &mut Vec<S>) -> bool, first: impl Fn(usize, T) -> (S, bool) + Copy,
    step: impl Fn(S, T) -> (S, bool) + Copy) -> bool {
    if data.is_empty() { return true; }
    if inner == 1 && len >= SHORT_LANE { return data.chunks(len).fold(true, |fits, l| lane(l, out) & fits); }
    if inner == 1 {
        return unsafe {
            appended(out, data.len(), |slots| {
                let whole = data.len() / (LANES * len) * (LANES * len);
                let (groups, rest) = slots.split_at_mut(whole);
                // SAFETY: `side_by_side` gives each position of the data it runs on, which holds as many items as the slots written here.
                side_by_side::<LANES, _, _>(&data[..whole], len, |x| first(0, x), step, |at, s| {
                    groups.get_unchecked_mut(at).write(s);
                }, |_, _| {}) & side_by_side::<1, _, _>(&data[whole..], len, |x| first(0, x), step, |at, s| {
                    rest.get_unchecked_mut(at).write(s);
                }, |_, _| {})
            })
        };
    }
    // Each row steps from the row before it, which `out` already holds.
    unsafe {
        appended(out, data.len(), |slots| {
            let mut fits = true;
            let mut at = 0;
            while at < data.len() {
                let (row, outs) = (&data[at..at + inner], &mut slots[at..at + inner]);
                for k in 0..inner {
                    let (t, f) = first(k, row[k]);
                    outs[k].write(t);
                    fits &= f;
                }
                let mut r = at + inner;
                while r < at + len * inner {
                    let (done, rest) = slots.split_at_mut(r);
                    let (before, row, outs) = (written(&mut done[r - inner..]), &data[r..r + inner], &mut rest[..inner]);
                    for k in 0..inner {
                        let (t, f) = step(before[k], row[k]);
                        outs[k].write(t);
                        fits &= f;
                    }
                    r += inner;
                }
                at += len * inner;
            }
            fits
        })
    }
}

/// A function of two Booleans, held as its truth table: `Logical(t)` gives `t[r][x]` for the running result `r` and the next item `x`. Its scans run any such function.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Logical(pub [[bool; 2]; 2]);

impl Logical {
    /// `r f x`, with no branch.
    #[inline(always)]
    fn apply(self, r: bool, x: bool) -> bool { let [[f00, f01], [f10, f11]] = self.0; (r & ((x & f11) | (!x & f10))) | (!r & ((x & f01) | (!x & f00))) }

    /// Appends the running results of `x` to `out`: the first item, or `start` combined with it where `start` is given, then each result combined with the next item.
    ///
    /// For each value of the next item, the function maps the running result to itself, to its negation, or to a constant. Where one value maps it to itself and the other to a constant, as for `and` and `or`, the results change at most once. `first_where` finds that place, and two runs fill the rest. Otherwise each map is `r ↦ (s ∧ r) ≠ t` for two bits `s` and `t`, and the maps compose as `r ↦ (s₂ ∧ s₁ ∧ r) ≠ ((s₂ ∧ t₁) ≠ t₂)`. Each block of 64 items packs into a word of bits, whose maps compose in six steps of shifts, with no branch and no step for each item.
    pub fn scan(self, x: &[bool], start: Option<bool>, out: &mut Vec<bool>) {
        let Some((&head, rest)) = x.split_first() else { return };
        let first = self.first(start, head);
        out.push(first);
        // `maps[v]` gives the results of an item `v` for a running result of false and of true.
        let [[f00, f01], [f10, f11]] = self.0;
        let maps = [[f00, f10], [f01, f11]];
        for v in [false, true] {
            let [lo, hi] = maps[usize::from(!v)];
            if maps[usize::from(v)] == [false, true] && lo == hi {
                let at = first_where(rest, |x| x != v);
                out.extend(std::iter::repeat_n(first, at));
                out.extend(std::iter::repeat_n(lo, rest.len() - at));
                return;
            }
        }
        let bit_maps = self.bit_maps();
        mapped(rest.len(), first, |start| bit_maps(packed(rest, start)), out);
    }

    /// A lane's first result: its first item `x`, combined with `start` where it's given.
    #[inline(always)]
    fn first(self, start: Option<bool>, x: bool) -> bool { start.map_or(x, |s| self.apply(s, x)) }

    /// The maps of 64 items given as the bits of a word, as [`mapped`] takes them: item `k` maps `r` to `(s ∧ r) ≠ t` for bit `k` of each word of `(s, t)`.
    #[inline(always)]
    fn bit_maps(self) -> impl Fn(u64) -> (u64, u64) + Copy {
        let [[f00, f01], [f10, f11]] = self.0;
        let word = |b: bool| if b { u64::MAX } else { 0 };
        // An item `v` maps a result of false to `f0v` and a result of true to `f1v`.
        let (s0, t0, s1, t1) = (word(f00 != f10), word(f00), word(f01 != f11), word(f01));
        move |bits| ((bits & s1) | (!bits & s0), (bits & t1) | (!bits & t0))
    }

    /// Appends the running results of `data`'s lanes of `len` consecutive items, 1 to 64, each lane from `start`, to `out` in one [`mapped`] call. The Boolean reductions' [`Reduction::short_lanes`] runs through it. Each word restarts the lanes that start in it with constant maps, at places that shift by `64 % len` from one word to the next. A lane's first result depends only on its first item, so it comes from the same bits.
    pub(super) fn short_lanes(self, data: &[bool], len: usize, start: Option<bool>, out: &mut Vec<bool>) {
        let (bit_maps, word) = (self.bit_maps(), |b: bool| if b { u64::MAX } else { 0 });
        let (yes, no) = (word(self.first(start, true)), word(self.first(start, false)));
        let starts = (0..64).step_by(len).fold(0u64, |p, k| p | 1 << k);
        let (shift, at) = (64 % len, std::cell::Cell::new(0));
        mapped(data.len(), false, |from| {
            let (bits, k) = (packed(data, from), at.get());
            at.set(if k >= shift { k - shift } else { k + len - shift });
            let (places, (s, t)) = (starts << k, bit_maps(bits));
            (s & !places, (t & !places) | ((bits & yes) | (!bits & no)) & places)
        }, out);
    }
}

/// The 64 Booleans of `block` as the bits of a word, the first in the lowest bit. Each 8 Booleans, one byte each, gather into a byte through one multiplication.
#[inline(always)]
fn pack(block: &[bool; 64]) -> u64 {
    let mut bits = 0;
    for (k, b) in block.as_chunks::<8>().0.iter().enumerate() {
        let bytes = (0..8).fold(0u64, |w, i| w | u64::from(b[i]) << (8 * i));
        bits |= (bytes.wrapping_mul(0x0102_0408_1020_4080) >> 56) << (8 * k);
    }
    bits
}

/// Expands the 64 bits of `r` into Boolean bytes, the lowest bit first, then converts that block into `out`. Each byte of `r` becomes 8 bytes of 0 or 1, bit `j` in byte `j`, through one multiplication, two masks, an add and a shift that's the same for every byte. LLVM compiles a test of bit `j` in byte `j` to a different shift for each byte, which wasm's SIMD can only emulate.
#[inline(always)]
fn unpack<S>(r: u64, out: &mut [std::mem::MaybeUninit<S>; 64]) where bool: Source<S> {
    let mut block = [false; 64];
    for (k, lanes) in block.as_chunks_mut::<8>().0.iter_mut().enumerate() {
        let bits = (r >> (8 * k) & 0xff).wrapping_mul(0x0101_0101_0101_0101) & 0x8040_2010_0804_0201;
        let bytes = ((bits + 0x7f7f_7f7f_7f7f_7f7f) >> 7 & 0x0101_0101_0101_0101).to_le_bytes();
        *lanes = std::array::from_fn(|j| bytes[j] != 0);
    }
    crate::slice::pairs_into(Arg::Many(&block), Arg::One(()), 64, out, |b, ()| (b.read(), true));
}

/// The results of 64 items, bit by bit, from the running result `carry`, where item `k`'s map is `r ↦ (s ∧ r) ≠ t` for bit `k` of `s` and `t`. Each step composes each bit's map with the map it holds from the bits before it. A bit with none before it takes the identity map, whose `s` is 1 and whose `t` is 0.
#[inline(always)]
fn composed(mut s: u64, mut t: u64, carry: bool) -> u64 {
    for k in [1, 2, 4, 8, 16, 32] {
        t ^= s & (t << k);
        s &= (s << k) | ((1u64 << k) - 1);
    }
    (s & if carry { u64::MAX } else { 0 }) ^ t
}

/// The Booleans of `x` from `start`, at most 64, as the bits of a word, the first in the lowest bit. Bits past the end of `x` are 0.
#[inline(always)]
fn packed(x: &[bool], start: usize) -> u64 {
    match x[start..].first_chunk::<64>() {
        Some(block) => pack(block),
        None => {
            let mut block = [false; 64];
            block[..x.len() - start].copy_from_slice(&x[start..]);
            pack(&block)
        }
    }
}

/// Appends the running results of `n` items to `out`, from the running result `carry` before the first. `maps(start)` gives the maps of the 64 items from `start` as two words of bits `(s, t)`: item `start + k` maps `r` to `(s ∧ r) ≠ t` for bit `k` of each. Bits past the last item are ignored.
#[inline(always)]
fn mapped<S: Copy>(n: usize, mut carry: bool, maps: impl Fn(usize) -> (u64, u64), out: &mut Vec<S>) where bool: Source<S> {
    unsafe {
        appended(out, n, |slots| {
            let (outs, last) = slots.as_chunks_mut::<64>();
            for (i, o) in outs.iter_mut().enumerate() {
                let (s, t) = maps(64 * i);
                let r = composed(s, t, carry);
                unpack(r, o);
                carry = r >> 63 != 0;
            }
            if !last.is_empty() {
                let (s, t) = maps(n - last.len());
                let r = composed(s, t, carry);
                for (k, slot) in last.iter_mut().enumerate() { slot.write(Source::<S>::read(r >> k & 1 != 0)); }
            }
        })
    };
}

/// Appends a Boolean-valued scan of long consecutive lanes, keeping an unseeded lane's first item as it is. Comparisons at carries 0 and 1 make independent two-state maps, whose packed words compose through `mapped`. A failed speculative comparison leaves `out` unchanged and declines, so the ordered runner checks only steps the scan actually takes.
pub fn comparison_scan<T: Element + PartialEq, S: Element>(
    data: &[T], len: usize, start: Option<T>, step: impl Fn(T, T) -> (T, Option<Kind>) + Copy, read: impl Fn(T) -> S + Copy,
    bulk: impl Fn(T, &[T], &mut Vec<bool>) -> Option<bool> + Copy, out: &mut Vec<S>) -> bool where bool: Source<T> + Source<S> {
    if len < SHORT_LANE { return false; }
    let (zero, one, before) = (false.read(), true.read(), out.len());
    let (mut low, mut high) = (Vec::new(), Vec::new());
    for lane in data.chunks(len) {
        let (first, used) = match start { Some(seed) => (step(seed, lane[0]), 1), None => { out.push(read(lane[0])); (step(lane[0], lane[1]), 2) } };
        let (first, failed) = first;
        if failed.is_some() { out.truncate(before); return false; }
        out.push(read(first));
        let rest = &lane[used..];
        low.clear();
        high.clear();
        let transitions = |carry, dest: &mut Vec<bool>| {
            if let Some(valid) = bulk(carry, rest, dest) { return valid; }
            unsafe {
                appended(dest, rest.len(), |slots| {
                    let mut valid = true;
                    for (&x, slot) in rest.iter().zip(slots) {
                        let (value, failed) = step(carry, x);
                        slot.write(value != zero);
                        valid &= failed.is_none();
                    }
                    valid
                })
            }
        };
        if !(transitions(zero, &mut low) & transitions(one, &mut high)) { out.truncate(before); return false; }
        mapped(rest.len(), first != zero, |at| { let (l, h) = (packed(&low, at), packed(&high, at)); (l ^ h, l) }, out);
    }
    true
}

/// Appends each cell's left difference scan without materializing its negated items. The first item is unchanged without a seed; every later item is subtracted from the running total.
pub(super) fn negated_cells<T: Element, R: Reduction<T, Total = T>>(
    items: &[T], dims: [usize; 3], seed: Option<&[T]>, out: &mut Vec<T>, neg: impl Fn(T) -> T + Copy) {
    let start = |i| seed.map(|s| s[if s.len() == 1 { 0 } else { i }]);
    let lane = |lane: &[T], out: &mut Vec<T>| {
        let (total, rest) = match start(0) { Some(s) => (s, lane), None => { out.push(lane[0]); (lane[0], &lane[1..]) } };
        grouped(rest, Some(total), out, neg, R::step, |x| x, T::DTYPE.is_integral());
        true
    };
    rows_scan(items, dims, out, lane, |i, x| (start(i).map_or(x, |s| R::step(s, neg(x))), true), |t, x| (R::step(t, neg(x)), true));
}