1use crate::error::{Error, Kind};
3pub use crate::slice::running;
4use crate::{
5 Arg, Float, Native, agreement::{Mapping, Read, Walked, expanded_items}, arith::Overflowing, buffer::{Buffer, Cast, DType, Element, Source, View, Width},
6 float::sealed::{Core, Folds}, gather::Steps, layout::Lanes, linalg::Scalar, parallel, slice::{CHUNK, appended, written}
7};
8use half::f16;
9use num_complex::Complex;
10use std::ops::Range;
11
12#[derive(Clone, Copy, Debug, PartialEq, Eq)]
13pub enum Direction { Left, Right }
14
15#[derive(Clone, Copy)]
16pub enum Seed<'a> { Identity, Initial(crate::agreement::Read<'a>), Tail(crate::agreement::Read<'a>) }
17impl<'a> Seed<'a> {
18 pub fn read(self) -> Option<crate::agreement::Read<'a>> { match self { Self::Initial(read) | Self::Tail(read) => Some(read), Self::Identity => None } }
20}
21
22#[derive(Clone, Copy, Debug)]
24pub struct Total { pub items: usize, pub bits: u64 }
25impl Total {
26 pub fn value<T: crate::buffer::Cast>(self) -> Result<T, Error> {
27 use crate::buffer::{Floats, Ints};
28 match T::DTYPE {
29 DType::Boolean => T::cast(View::Booleans(&[self.bits != 0]), 0),
30 DType::Integer(width) => {
31 let value = match width {
32 Width::U8 => self.bits as u8 as i64, Width::I16 => self.bits as i16 as i64, Width::I32 => self.bits as i32 as i64,
33 Width::I64 => self.bits as i64
34 };
35 T::cast(View::Integers(Ints::I64(&[value])), 0)
36 }
37 DType::Float(_) => T::cast(View::Floats(Floats::F64(&[f64::from_bits(self.bits)])), 0),
38 _ => Err(Error::new(Kind::Domain, "compiled totals have a real machine dtype")),
39 }
40 }
41}
42
43pub enum Piece<'a> { Items(View<'a>), Total(Total) }
44
45pub type Pieces<'a> = dyn Fn(std::ops::Range<usize>, &mut dyn FnMut(usize, Piece<'_>)) -> Result<(), Error> + Sync + 'a;
46
47#[inline]
49pub fn sum<T: Copy>(x: &[T]) -> (<Sum as Reduction<T>>::Total, bool) where Sum: Reduction<T> { Sum::slice(x) }
50#[inline]
52pub fn prod<T: Copy>(x: &[T]) -> (<Product as Reduction<T>>::Total, bool) where Product: Reduction<T> { Product::slice(x) }
53#[inline]
55pub fn max<T: Copy>(x: &[T]) -> T where Max: Reduction<T, Total = T> { Max::slice(x).0 }
56#[inline]
58pub fn min<T: Copy>(x: &[T]) -> T where Min: Reduction<T, Total = T> { Min::slice(x).0 }
59#[inline]
61pub fn extremes<T: Copy>(x: &[T]) -> (T, T) where Min: Reduction<T, Total = T>, Max: Reduction<T, Total = T> { extremes_of(x, |v| v) }
62
63#[inline]
65pub(crate) fn extremes_of<A: Copy, T: Copy>(x: &[A], key: impl Fn(A) -> T) -> (T, T) where Min: Reduction<T, Total = T>, Max: Reduction<T, Total = T> {
66 let (min, max) = (<Min as Reduction<T>>::step, <Max as Reduction<T>>::step);
67 x.iter().fold((<Min as Reduction<T>>::unit(), <Max as Reduction<T>>::unit()), |(lo, hi), &a| {
68 let v = key(a);
69 (min(lo, v), max(hi, v))
70 })
71}
72
73#[inline(always)]
75pub fn maximum<F: Float>(x: F, y: F) -> F { F::from_working(Core::maximum(x.working(), y.working())) }
76
77#[inline(always)]
79pub fn minimum<F: Float>(x: F, y: F) -> F { F::from_working(Core::minimum(x.working(), y.working())) }
80
81#[derive(Clone, Copy, Debug)]
83pub struct Maximum;
84forms!(#[inline] Maximum at real: |_, [x, y]| crate::slice::Total(maximum(x, y)));
85forms!(#[inline] Maximum at u8, i16, i32, i64: |_, [x, y]| x.max(y));
86forms!(#[inline] Maximum at bool: |_, [x, y]| x | y);
87
88#[derive(Clone, Copy, Debug)]
90pub struct Minimum;
91forms!(#[inline] Minimum at real: |_, [x, y]| crate::slice::Total(minimum(x, y)));
92forms!(#[inline] Minimum at u8, i16, i32, i64: |_, [x, y]| x.min(y));
93forms!(#[inline] Minimum at bool: |_, [x, y]| x & y);
94
95pub trait Reduction<T> {
97 type Total: Copy + Send + Sync;
99 const COMMUTATIVE: bool = false;
101 fn unit() -> Self::Total;
103 #[inline(always)]
105 fn identity() -> Self::Total { Self::unit() }
106 fn step(acc: Self::Total, x: T) -> Self::Total;
108 #[inline(always)]
110 fn checked(acc: Self::Total, x: T) -> (Self::Total, bool) { (Self::step(acc, x), true) }
111 #[inline(always)]
113 fn one(x: T) -> (Self::Total, bool) { Self::checked(Self::identity(), x) }
114 fn merge(a: Self::Total, b: Self::Total) -> (Self::Total, bool);
116 #[inline(always)]
118 fn chain(a: Self::Total, b: Self::Total) -> (Self::Total, bool) { Self::merge(a, b) }
119 fn slice(x: &[T]) -> (Self::Total, bool) where Self: Sized, T: Copy {
121 if x.is_empty() { return (Self::unit(), true); }
122 x.iter().fold((Self::identity(), true), |(t, fits), &v| { let (t, ok) = Self::checked(t, v); (t, fits & ok) })
123 }
124 fn scan(x: &[T], start: Option<Self::Total>, out: &mut Vec<Self::Total>) -> bool where Self: Sized, T: Copy { checked_steps::<T, Self>(x, start, out) }
126 #[inline]
128 fn rows(data: &[T], columns: usize, out: &mut [Self::Total]) -> bool where Self: Sized, T: Copy { fold_rows::<T, Self>(data, columns, out) }
129 #[inline]
131 fn short_lanes(_data: &[T], _len: usize, _start: Option<Self::Total>, _out: &mut Vec<Self::Total>) -> Option<bool> where Self: Sized { None }
132 const SPLIT: usize = 8;
134}
135
136pub trait Totals: Copy {
138 fn buffer_of(totals: Vec<Self>) -> crate::buffer::Buffer;
140 fn from_partial(_: Total) -> Result<Self, Error> { Err(Error::new(Kind::Domain, "this reduction cannot read compiled partial totals")) }
141}
142impl<T: crate::buffer::Cast> Totals for T {
143 #[inline]
144 fn buffer_of(totals: Vec<T>) -> crate::buffer::Buffer { T::buffer(totals) }
145 fn from_partial(total: Total) -> Result<Self, Error> { total.value() }
146}
147impl<S: Totals> Totals for (S, bool) {
148 #[inline]
149 fn buffer_of(totals: Vec<(S, bool)>) -> crate::buffer::Buffer { S::buffer_of(totals.into_iter().map(|(s, _)| s).collect()) }
150}
151impl<F: Native + Element> Totals for Scaled<F> {
152 #[inline]
153 fn buffer_of(totals: Vec<Scaled<F>>) -> crate::buffer::Buffer { F::buffer(totals.into_iter().map(Scaled::value).collect()) }
154}
155
156#[derive(Clone, Copy, Debug)]
158pub struct Sum;
159#[derive(Clone, Copy, Debug)]
161pub struct Product;
162#[derive(Clone, Copy, Debug)]
164pub struct Max;
165#[derive(Clone, Copy, Debug)]
167pub struct Min;
168#[derive(Clone, Copy, Debug)]
170pub struct And;
171#[derive(Clone, Copy, Debug)]
173pub struct Or;
174#[derive(Clone, Copy, Debug)]
176pub struct Xor;
177#[derive(Clone, Copy, Debug)]
179pub struct Xnor;
180forms!(Xor at bool: |_, [x, y]| x != y);
181forms!(Xnor at bool: |_, [x, y]| x == y);
182
183#[inline(always)]
185pub(crate) fn folded<A: Copy, S: Copy>(x: &[A], unit: S, add: impl Fn(S, A) -> S + Copy, merge: impl Fn(S, S) -> S) -> S {
186 let merge = &merge;
187 running_results(x, unit, add, merge, move |parts, rest| rest.iter().fold(parts.iter().fold(unit, |s, &p| merge(s, p)), |s, &v| add(s, v)))
188}
189
190#[inline(always)]
192fn running_results<A: Copy, S: Copy, T>(x: &[A], unit: S, add: impl Fn(S, A) -> S, merge: impl Fn(S, S) -> S, finish: impl FnOnce(&[S], &[A]) -> T) -> T {
193 match crate::slice::running(size_of::<S>()).min(x.len() / 8) {
194 ..8 => running_by::<A, S, T, 4>(x, unit, add, merge, finish), 8..16 => running_by::<A, S, T, 8>(x, unit, add, merge, finish),
195 16..32 => running_by::<A, S, T, 16>(x, unit, add, merge, finish), _ => running_by::<A, S, T, 32>(x, unit, add, merge, finish)
196 }
197}
198
199#[inline(always)]
201fn running_by<A: Copy, S: Copy, T, const R: usize>(
202 x: &[A], unit: S, add: impl Fn(S, A) -> S, merge: impl Fn(S, S) -> S, finish: impl FnOnce(&[S], &[A]) -> T) -> T {
203 let (chunks, rest) = x.as_chunks::<R>();
204 let parts = if chunks.len() * R <= LEAF { leaf(chunks, unit, &add) } else { tree(chunks, unit, &add, &merge) };
205 finish(&parts, rest)
206}
207
208const LEAF: usize = 4096;
210
211fn tree<A: Copy, S: Copy, const R: usize>(chunks: &[[A; R]], unit: S, add: &impl Fn(S, A) -> S, merge: &impl Fn(S, S) -> S) -> [S; R] {
213 if chunks.len() * R > LEAF {
214 let (a, b) = chunks.split_at(chunks.len() / 2);
215 let (mut a, b) = (tree(a, unit, add, merge), tree(b, unit, add, merge));
216 for k in 0..R { a[k] = merge(a[k], b[k]) }
217 return a;
218 }
219 leaf(chunks, unit, add)
220}
221
222#[inline(always)]
223fn leaf<A: Copy, S: Copy, const R: usize>(chunks: &[[A; R]], unit: S, add: &impl Fn(S, A) -> S) -> [S; R] {
224 let mut parts = [unit; R];
225 for chunk in chunks { for k in 0..R { parts[k] = add(parts[k], chunk[k]) } }
226 parts
227}
228
229const SCAN_BLOCK: usize = 64;
231
232#[inline(always)]
234pub(crate) fn grouped<A: Copy, S: Copy, O>(
235 x: &[A], start: Option<S>, out: &mut Vec<O>, from: impl Fn(A) -> S, f: impl Fn(S, S) -> S, to: impl Fn(S) -> O, reorders: bool) {
236 let Some((&first, rest)) = x.split_first() else { return };
237 unsafe {
238 appended(out, x.len(), |slots| {
239 let first = from(first);
240 let mut total = start.map_or(first, |s| f(s, first));
241 slots[0].write(to(total));
242 let ((groups, tail), (outs, last)) = (rest.as_chunks::<4>(), slots[1..].as_chunks_mut::<4>());
243 let scanned = |[a, b, c, d]: [A; 4]| {
244 let (a, b, c, d) = (from(a), from(b), from(c), from(d));
245 let ab = f(a, b);
246 [a, ab, f(ab, c), f(ab, f(c, d))]
247 };
248 let (mut sums, mut ends) = ([std::mem::MaybeUninit::<[S; 4]>::uninit(); SCAN_BLOCK], [std::mem::MaybeUninit::<S>::uninit(); SCAN_BLOCK]);
249 let mut i0 = 0;
251 while i0 < groups.len() {
252 let n = (groups.len() - i0).min(SCAN_BLOCK);
253 let (groups, outs) = (&groups[i0..i0 + n], &mut outs[i0..i0 + n]);
254 if reorders {
255 for g in 0..n {
256 let s = scanned(groups[g]);
257 sums[g].write(s);
258 ends[g].write(s[3]);
259 }
260 }
261 for g in 0..n {
262 let sums = if reorders { sums[g].assume_init() } else { scanned(groups[g]) };
264 for k in 0..4 { outs[g][k].write(to(f(total, sums[k]))); }
266 total = f(total, if reorders { ends[g].assume_init() } else { sums[3] });
268 }
269 i0 += n;
270 }
271 for (o, &v) in last.iter_mut().zip(tail) {
272 total = f(total, from(v));
273 o.write(to(total));
274 }
275 })
276 }
277}
278
279#[inline]
281fn first_where<T: Copy>(x: &[T], test: impl Fn(T) -> bool) -> usize {
282 let (blocks, _) = x.as_chunks::<64>();
283 let from = blocks.iter().position(|b| b.iter().fold(false, |any, &v| any | test(v))).unwrap_or(blocks.len()) * 64;
284 x[from..].iter().position(|&v| test(v)).map_or(x.len(), |i| from + i)
285}
286
287#[inline(always)]
289fn moved<T: PartialEq>(pick: usize, before: T, after: T, k: usize) -> usize { if after != before { k } else { pick } }
290
291macro_rules! floats {
293 ($($r:ident: $unit:ident, $(identity: $identity:expr,)? $(rows: $rows:ident,)? $(chain: $chain:expr,)? |$a:ident, $b:ident| $step:expr, $fold:expr, $scan:expr;)+) => {$(
294 impl<F: Native> Reduction<F> for $r {
295 type Total = F;
296 const COMMUTATIVE: bool = true;
297 const SPLIT: usize = 64;
298 #[inline(always)]
299 fn unit() -> F { F::$unit() }
300 #[inline(always)]
301 fn step($a: F, $b: F) -> F { $step }
302 #[inline(always)]
303 fn merge(a: F, b: F) -> (F, bool) { (Self::step(a, b), true) }
304 $(#[inline(always)]
305 fn identity() -> F { $identity })?
306 $(#[inline]
307 fn rows(data: &[F], columns: usize, out: &mut [F]) -> bool { $rows::<F, Self>(data, columns, out) })?
308 $(#[inline(always)]
309 fn chain(a: F, b: F) -> (F, bool) { (($chain)(a, b), true) })?
310 #[inline]
311 fn slice(x: &[F]) -> (F, bool) { (($fold)(x), true) }
312 #[inline]
313 fn scan(x: &[F], start: Option<F>, out: &mut Vec<F>) -> bool {
314 ($scan)(x, start, out);
315 true
316 }
317 }
318 )+};
319}
320floats! {
321 Sum: zero, identity: -F::zero(), rows: rows_side_by_side, |a, b| a.algebraic_add(b), |x: &[F]| if x.is_empty() { <Self as Reduction<F>>::unit() } else { folded(x, <Self as Reduction<F>>::identity(), |a: F, b: F| a + b, |a: F, b: F| a + b) }, F::scan_sum;
322 Product: one, rows: rows_side_by_side, |a, b| a.algebraic_mul(b), |x| folded(x, F::one(), |a: F, b: F| a * b, |a: F, b: F| a * b), |x, start, out| grouped(x, start, out, |v| v, |a: F, b: F| a * b, |t| t, false);
323 Max: neg_infinity, chain: F::step_max, |acc, x| x.maximum(acc), F::fold_max, F::scan_max;
324 Min: infinity, chain: F::step_min, |acc, x| x.minimum(acc), F::fold_min, F::scan_min;
325}
326
327macro_rules! halves {
329 ($total:ty, $wide:expr, $narrow:expr; $($r:ty: $add:expr),+) => {$(
330 impl Reduction<f16> for $r {
331 type Total = $total;
332 const COMMUTATIVE: bool = true;
333 #[inline(always)]
334 fn unit() -> $total { $narrow(<$r as Reduction<f32>>::unit()) }
335 #[inline(always)]
336 fn step(acc: $total, x: f16) -> $total { $narrow(<$r as Reduction<f32>>::step($wide(acc), x.to_f32())) }
337 #[inline(always)]
338 fn merge(a: $total, b: $total) -> ($total, bool) { ($narrow(<$r as Reduction<f32>>::merge($wide(a), $wide(b)).0), true) }
339 #[inline(always)]
340 fn identity() -> $total { $narrow(<$r as Reduction<f32>>::identity()) }
341 #[inline]
342 fn slice(x: &[f16]) -> ($total, bool) {
343 let (unit, add) = (<$r as Reduction<f32>>::identity(), $add);
344 ($narrow(if x.is_empty() { <$r as Reduction<f32>>::unit() } else { folded(x, unit, |s, v| add(s, v.to_f32()), add) }), true)
345 }
346 #[inline]
347 fn scan(x: &[f16], start: Option<$total>, out: &mut Vec<$total>) -> bool {
348 grouped(x, start.map($wide), out, |v: f16| <$r as Reduction<f32>>::one(v.to_f32()).0, <$r as Reduction<f32>>::step, $narrow, false);
349 true
350 }
351 }
352 )+};
353}
354halves!(f32, |t: f32| t, |t: f32| t; Sum: |a: f32, b: f32| a + b, Product: |a: f32, b: f32| a * b);
356halves!(f16, f16::to_f32, f16::from_f32; Max: <Max as Reduction<f32>>::step, Min: <Min as Reduction<f32>>::step);
357
358impl<F: Native> Reduction<Complex<F>> for Sum {
360 type Total = Complex<F>;
361 const COMMUTATIVE: bool = true;
362 #[inline(always)]
363 fn unit() -> Complex<F> { Complex::new(F::zero(), F::zero()) }
364 #[inline(always)]
365 fn step(acc: Complex<F>, x: Complex<F>) -> Complex<F> { Complex::new(acc.re.algebraic_add(x.re), acc.im.algebraic_add(x.im)) }
366 #[inline(always)]
367 fn merge(a: Complex<F>, b: Complex<F>) -> (Complex<F>, bool) { (<Self as Reduction<Complex<F>>>::step(a, b), true) }
368 #[inline(always)]
369 fn identity() -> Complex<F> { Complex::new(-F::zero(), -F::zero()) }
370 #[inline]
371 fn slice(x: &[Complex<F>]) -> (Complex<F>, bool) {
372 let add = |a: Complex<F>, b: Complex<F>| a + b;
373 (if x.is_empty() { <Self as Reduction<Complex<F>>>::unit() } else { folded(x, <Self as Reduction<Complex<F>>>::identity(), add, add) }, true)
374 }
375 #[inline]
376 fn scan(x: &[Complex<F>], start: Option<Complex<F>>, out: &mut Vec<Complex<F>>) -> bool {
377 grouped(x, start, out, |v| v, |a: Complex<F>, b| a + b, |t| t, false);
378 true
379 }
380}
381
382impl<F: Native> Reduction<Complex<F>> for Product {
384 type Total = Complex<F>;
385 const COMMUTATIVE: bool = true;
386 #[inline(always)]
387 fn unit() -> Complex<F> { Complex::new(F::one(), F::zero()) }
388 #[inline(always)]
389 fn step(acc: Complex<F>, x: Complex<F>) -> Complex<F> { acc * x }
390 #[inline(always)]
391 fn merge(a: Complex<F>, b: Complex<F>) -> (Complex<F>, bool) { (a * b, true) }
392 #[inline]
393 fn slice(x: &[Complex<F>]) -> (Complex<F>, bool) {
394 let step = <Self as Reduction<Complex<F>>>::step;
395 (folded(x, <Self as Reduction<Complex<F>>>::unit(), step, step), true)
396 }
397 #[inline]
398 fn scan(x: &[Complex<F>], start: Option<Complex<F>>, out: &mut Vec<Complex<F>>) -> bool {
399 grouped(x, start, out, |v| v, |a: Complex<F>, b| a * b, |t| t, false);
400 true
401 }
402}
403
404#[cfg(any(target_arch = "aarch64", all(target_arch = "wasm32", target_feature = "simd128")))]
406mod lanes {
407 use super::*;
408 #[cfg(target_arch = "wasm32")]
409 use core::arch::wasm32::*;
410 #[cfg(target_arch = "aarch64")]
411 use core::arch::{aarch64::*, asm};
412 const PARTS: usize = 8;
414 macro_rules! lanes {
415 ($($t:ty: $lanes:literal, $splat:ident, $load:ident, $max:ident, $min:ident, $step:ident($arg:tt), $scan:ident, $sum:ident;)+) => {$(
416 impl Folds for $t {
417 #[inline(always)]
418 fn step_max(t: $t, v: $t) -> $t { lanes!(@$step $arg, "fmaxnm", $splat, $max, t, v) }
419 #[inline(always)]
420 fn step_min(t: $t, v: $t) -> $t { lanes!(@$step $arg, "fminnm", $splat, $min, t, v) }
421 fn fold_max(x: &[$t]) -> $t { lanes!(@fold x, $t, $lanes, NEG_INFINITY, $splat, $load, $max, Self::step_max) }
422 fn fold_min(x: &[$t]) -> $t { lanes!(@fold x, $t, $lanes, INFINITY, $splat, $load, $min, Self::step_min) }
423 fn scan_max(x: &[$t], start: Option<$t>, out: &mut Vec<$t>) {
424 lanes!(@scan $scan, x, start, out, lanes!(@read $step, $t, NEG_INFINITY), Self::step_max, below, |b: &[$t], t: $t| Self::fold_max(b) < t)
425 }
426 fn scan_min(x: &[$t], start: Option<$t>, out: &mut Vec<$t>) {
427 lanes!(@scan $scan, x, start, out, lanes!(@read $step, $t, INFINITY), Self::step_min, above, |b: &[$t], t: $t| Self::fold_min(b) > t)
428 }
429 fn scan_sum(x: &[$t], start: Option<$t>, out: &mut Vec<$t>) { lanes!(@sum $sum, x, start, out) }
430 #[inline(always)]
431 fn ordered_lanes<const W: usize>(a: [$t; W], c: [$t; W]) -> ([$t; W], [$t; W]) {
432 const { assert!(W == $lanes) };
433 unsafe {
435 let (x, y) = ($load(a.as_ptr().cast()), $load(c.as_ptr().cast()));
436 (core::mem::transmute_copy(&$min(x, y)), core::mem::transmute_copy(&$max(y, x)))
437 }
438 }
439 }
440 )+};
441 (@fold $x:ident, $t:ty, $lanes:literal, $unit:ident, $splat:ident, $load:ident, $op:ident, $step:expr) => {{
442 let (chunks, rest) = $x.as_chunks::<{ PARTS * $lanes }>();
443 let lanes: [$t; $lanes] = unsafe {
445 let mut parts = [$splat(<$t>::$unit); PARTS];
446 for c in chunks { for k in 0..PARTS { parts[k] = $op(parts[k], $load(c.as_ptr().add(k * $lanes).cast())) } }
447 core::mem::transmute(parts.into_iter().reduce(|a, b| $op(a, b)).expect("a fold keeps `PARTS` vectors"))
448 };
449 lanes.into_iter().chain(rest.iter().copied()).fold(<$t>::$unit, $step)
450 }};
451 (@asm $reg:literal, $ins:literal, $splat:ident, $op:ident, $x:ident, $y:ident) => {{
452 let r;
453 unsafe { asm!(concat!($ins, " {0:", $reg, "}, {1:", $reg, "}, {2:", $reg, "}"), lateout(vreg) r, in(vreg) $x, in(vreg) $y, options(pure, nomem, nostack)) }
455 r
456 }};
457 (@lane $lane:ident, $ins:literal, $splat:ident, $op:ident, $x:ident, $y:ident) => { $lane::<0>($op($splat($x), $splat($y))) };
458 (@read asm, $t:ty, $unit:ident) => { |v: $t| v };
460 (@read lane, $t:ty, $unit:ident) => { |v: $t| if v.is_nan() { <$t>::$unit } else { v } };
461 (@scan grouped, $x:ident, $start:ident, $out:ident, $from:expr, $step:expr, $compare:expr, $fold:expr) => { grouped($x, $start, $out, $from, $step, |t| t, false) };
462 (@scan compare, $x:ident, $start:ident, $out:ident, $from:expr, $step:expr, $compare:expr, $fold:expr) => { records($x, $start, $out, $from, $step, $compare) };
463 (@scan fold, $x:ident, $start:ident, $out:ident, $from:expr, $step:expr, $compare:expr, $fold:expr) => { records($x, $start, $out, $from, $step, $fold) };
464 (@sum grouped, $x:ident, $start:ident, $out:ident) => { grouped($x, $start, $out, |v| v, |a, b| a + b, |t| t, false) };
465 (@sum vector, $x:ident, $start:ident, $out:ident) => { sum_scan($x, $start, $out) };
466 }
467 #[cfg(target_arch = "aarch64")]
468 lanes! {
469 f64: 2, vdupq_n_f64, vld1q_f64, vmaxnmq_f64, vminnmq_f64, asm("d"), grouped, grouped;
470 f32: 4, vdupq_n_f32, vld1q_f32, vmaxnmq_f32, vminnmq_f32, asm("s"), compare, vector;
471 }
472 #[cfg(target_arch = "aarch64")]
474 mod quad {
475 use core::arch::aarch64::*;
476 pub type V = float32x4_t;
477 #[inline(always)]
479 pub fn splat(x: f32) -> V { unsafe { vdupq_n_f32(x) } }
480 #[inline(always)]
481 pub unsafe fn load(p: *const f32) -> V { unsafe { vld1q_f32(p) } }
482 #[inline(always)]
483 pub unsafe fn store(p: *mut f32, v: V) { unsafe { vst1q_f32(p, v) } }
484 #[inline(always)]
485 pub fn add(a: V, b: V) -> V { unsafe { vaddq_f32(a, b) } }
486 #[inline(always)]
487 pub fn up1(fill: V, v: V) -> V { unsafe { vextq_f32::<3>(fill, v) } }
488 #[inline(always)]
489 pub fn up2(fill: V, v: V) -> V { unsafe { vextq_f32::<2>(fill, v) } }
490 #[inline(always)]
491 pub fn last(v: V) -> V { unsafe { vdupq_laneq_f32::<3>(v) } }
492 #[inline(always)]
493 pub fn first(v: V) -> f32 { unsafe { vgetq_lane_f32::<0>(v) } }
494 }
495 #[cfg(target_arch = "wasm32")]
496 mod quad {
497 use core::arch::wasm32::*;
498 pub type V = v128;
499 #[inline(always)]
500 pub fn splat(x: f32) -> V { f32x4_splat(x) }
501 #[inline(always)]
502 pub unsafe fn load(p: *const f32) -> V { unsafe { v128_load(p.cast()) } }
503 #[inline(always)]
504 pub unsafe fn store(p: *mut f32, v: V) { unsafe { v128_store(p.cast(), v) } }
505 #[inline(always)]
506 pub fn add(a: V, b: V) -> V { f32x4_add(a, b) }
507 #[inline(always)]
508 pub fn up1(fill: V, v: V) -> V { i32x4_shuffle::<3, 4, 5, 6>(fill, v) }
509 #[inline(always)]
510 pub fn up2(fill: V, v: V) -> V { i32x4_shuffle::<2, 3, 4, 5>(fill, v) }
511 #[inline(always)]
512 pub fn last(v: V) -> V { i32x4_shuffle::<3, 3, 3, 3>(v, v) }
513 #[inline(always)]
514 pub fn first(v: V) -> f32 { f32x4_extract_lane::<0>(v) }
515 }
516 fn sum_scan(x: &[f32], start: Option<f32>, out: &mut Vec<f32>) {
518 use quad::*;
519 unsafe {
520 appended(out, x.len(), |slots| {
521 let ((chunks, rest), (outs, tail)) = (x.as_chunks::<16>(), slots.as_chunks_mut::<16>());
522 let (fill, mut c) = (splat(-0.0), splat(start.unwrap_or(-0.0)));
523 for (o, x) in outs.iter_mut().zip(chunks) {
524 let scanned = |k: usize| {
525 let v = unsafe { load(x.as_ptr().add(4 * k)) };
527 let v = add(v, up1(fill, v));
528 add(v, up2(fill, v))
529 };
530 let (p0, p1, p2, p3) = (scanned(0), scanned(1), scanned(2), scanned(3));
531 let s0 = last(p0);
532 let s01 = add(s0, last(p1));
533 let s012 = add(s01, last(p2));
534 let o = o.as_mut_ptr().cast::<f32>();
535 unsafe {
537 store(o, add(c, p0));
538 store(o.add(4), add(add(c, s0), p1));
539 store(o.add(8), add(add(c, s01), p2));
540 store(o.add(12), add(add(c, s012), p3));
541 }
542 c = add(c, add(s012, last(p3)));
543 }
544 let mut total = first(c);
545 for (o, &v) in tail.iter_mut().zip(rest) {
546 total += v;
547 o.write(total);
548 }
549 })
550 }
551 }
552 #[cfg(target_arch = "wasm32")]
553 lanes! {
554 f64: 2, f64x2_splat, v128_load, f64x2_pmax, f64x2_pmin, lane(f64x2_extract_lane), fold, grouped;
555 f32: 4, f32x4_splat, v128_load, f32x4_pmax, f32x4_pmin, lane(f32x4_extract_lane), fold, vector;
556 }
557}
558
559#[cfg(not(any(target_arch = "aarch64", all(target_arch = "wasm32", target_feature = "simd128"))))]
561mod plain {
562 #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
563 use core::arch::x86_64::*;
564 macro_rules! plain {
565 ($($t:ty: $lanes:literal, $load:ident, $min:ident, $max:ident;)+) => {$(
566 impl super::Folds for $t {
567 #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
568 #[inline(always)]
569 fn ordered_lanes<const W: usize>(a: [$t; W], c: [$t; W]) -> ([$t; W], [$t; W]) {
570 const { assert!(W == $lanes) };
571 unsafe {
573 let (x, y) = ($load(a.as_ptr()), $load(c.as_ptr()));
574 (core::mem::transmute_copy(&$min(y, x)), core::mem::transmute_copy(&$max(x, y)))
575 }
576 }
577 }
578 )+};
579 }
580 plain! {
581 f32: 8, _mm256_loadu_ps, _mm256_min_ps, _mm256_max_ps;
582 f64: 4, _mm256_loadu_pd, _mm256_min_pd, _mm256_max_pd;
583 }
584}
585
586#[inline(always)]
588fn narrow_sum<T: Copy + Into<i64>>(x: &[T]) -> (i64, bool) {
589 let total: i128 = x.chunks(1 << 31).map(|c| i128::from(folded(c, 0i64, |s, v| s + v.into(), |a, b| a + b))).sum();
590 (total as i64, i64::try_from(total).is_ok())
591}
592
593#[inline(always)]
595fn small_sum<T: Copy + Into<i32>>(x: &[T]) -> (i64, bool) {
596 (x.chunks(1 << 15).map(|c| i64::from(folded(c, 0i32, |s, v| s + v.into(), |a, b| a + b))).sum(), true)
597}
598
599#[inline(always)]
601fn wide_sum(x: &[i64]) -> (i64, bool) {
602 let halves = |c: &[i64]| folded(c, (0i64, 0i64), |(h, l), v| (h + (v >> 32), l + (v & 0xFFFF_FFFF)), |(h, l), (g, m)| (h + g, l + m));
603 let total: i128 = x.chunks(1 << 31).map(|c| {
604 let (h, l) = halves(c);
605 (i128::from(h) << 32) + i128::from(l)
606 }).sum();
607 (total as i64, i64::try_from(total).is_ok())
608}
609
610#[inline(always)]
612fn signed_cumsum<T: Copy + Into<i64>>(x: &[T], start: i64, out: &mut Vec<i64>) -> bool {
613 unsafe {
614 appended(out, x.len(), |slots| {
615 let (mut total, mut overflow) = (start, 0i64);
616 for (o, &v) in slots.iter_mut().zip(x) {
617 let v: i64 = v.into();
618 let sum = total.wrapping_add(v);
619 overflow |= (total ^ sum) & (v ^ sum);
620 total = sum;
621 o.write(sum);
622 }
623 overflow >= 0
624 })
625 }
626}
627
628#[inline(always)]
630fn checked_scan<T: Copy + Into<S>, S: Copy>(x: &[T], start: S, out: &mut Vec<S>, op: impl Fn(S, S) -> (S, bool)) -> bool {
631 unsafe {
632 appended(out, x.len(), |slots| {
633 let (mut total, mut ok) = (start, true);
634 for (o, &v) in slots.iter_mut().zip(x) {
635 let (next, fits) = op(total, v.into());
636 (total, ok) = (next, ok & fits);
637 o.write(next);
638 }
639 ok
640 })
641 }
642}
643
644macro_rules! integers {
646 ($($t:ty: $sum:ident),+) => {$(
647 impl Reduction<$t> for Sum {
648 type Total = i64;
649 const COMMUTATIVE: bool = true;
650 #[inline(always)]
651 fn unit() -> i64 { 0 }
652 #[inline(always)]
653 fn step(acc: i64, x: $t) -> i64 { acc.plus(x.into()).0 }
654 #[inline(always)]
655 fn checked(acc: i64, x: $t) -> (i64, bool) { acc.plus(x.into()) }
656 #[inline(always)]
657 fn merge(a: i64, b: i64) -> (i64, bool) { a.plus(b) }
658 #[inline]
659 fn slice(x: &[$t]) -> (i64, bool) { $sum(x) }
660 #[inline]
661 fn scan(x: &[$t], start: Option<i64>, out: &mut Vec<i64>) -> bool { signed_cumsum(x, start.unwrap_or(0), out) }
662 }
663 impl Reduction<$t> for Product {
664 type Total = i64;
665 const COMMUTATIVE: bool = true;
666 #[inline(always)]
667 fn unit() -> i64 { 1 }
668 #[inline(always)]
669 fn step(acc: i64, x: $t) -> i64 { acc.times(x.into()).0 }
670 #[inline(always)]
671 fn checked(acc: i64, x: $t) -> (i64, bool) { acc.times(x.into()) }
672 #[inline(always)]
673 fn merge(a: i64, b: i64) -> (i64, bool) { a.times(b) }
674 #[inline]
675 fn slice(x: &[$t]) -> (i64, bool) {
676 x.iter().fold((1 as i64, true), |(p, ok), &v| { let (q, fits) = p.times(v.into()); (q, ok & fits) })
677 }
678 #[inline]
679 fn scan(x: &[$t], start: Option<i64>, out: &mut Vec<i64>) -> bool { checked_scan(x, start.unwrap_or(1), out, <i64>::times) }
680 }
681 )+};
682}
683integers!(u8: small_sum, i16: small_sum, i32: narrow_sum, i64: wide_sum);
684
685macro_rules! wrapped_sums {
687 ($($t:ty: $slice:expr),+) => {$(
688 impl Reduction<$t> for crate::Wrapped<Sum> {
689 type Total = i64;
690 const COMMUTATIVE: bool = true;
691 #[inline(always)]
692 fn unit() -> i64 { 0 }
693 #[inline(always)]
694 fn step(acc: i64, x: $t) -> i64 { acc.wrapping_add(x.into()) }
695 #[inline(always)]
696 fn merge(a: i64, b: i64) -> (i64, bool) { (a.wrapping_add(b), true) }
697 #[inline]
698 fn slice(x: &[$t]) -> (i64, bool) { (($slice)(x), true) }
699 #[inline]
700 fn scan(x: &[$t], start: Option<i64>, out: &mut Vec<i64>) -> bool { signed_cumsum(x, start.unwrap_or(0), out); true }
701 }
702 )+};
703}
704wrapped_sums!(u8: |x| small_sum(x).0, i16: |x| small_sum(x).0, i32: |x| narrow_sum(x).0, i64: |x: &[i64]| folded(x, 0i64, i64::wrapping_add, i64::wrapping_add));
705
706fn xnor_of<T: std::ops::Not<Output = T>>(xor: T, n: usize) -> T { if n % 2 == 0 { !xor } else { xor } }
708
709macro_rules! closed {
711 (@one $r:ident, $t:ty, $unit:expr, $step:expr, $scan:ident) => { closed!(@one $r, $t, $unit, $step, $scan, |x: &[$t]| x.iter().copied().fold($unit, $step)); };
712 (@one $r:ident, $t:ty, $unit:expr, $step:expr, $scan:ident, $slice:expr) => {
713 impl Reduction<$t> for $r {
714 type Total = $t;
715 const COMMUTATIVE: bool = true;
716 #[inline(always)]
717 fn unit() -> $t { $unit }
718 #[inline(always)]
719 fn step(acc: $t, x: $t) -> $t { ($step)(acc, x) }
720 #[inline(always)]
721 fn merge(a: $t, b: $t) -> ($t, bool) { (($step)(a, b), true) }
722 #[inline]
723 fn slice(x: &[$t]) -> ($t, bool) { (($slice)(x), true) }
724 #[inline]
725 fn scan(x: &[$t], start: Option<$t>, out: &mut Vec<$t>) -> bool {
726 closed!(@$scan x, start, out, $t, $step);
727 true
728 }
729 #[inline]
730 fn rows(data: &[$t], columns: usize, out: &mut [$t]) -> bool { rows_side_by_side::<$t, Self>(data, columns, out) }
731 }
732 };
733 (@records $x:ident, $start:ident, $out:ident, $t:ty, $step:expr) => { records($x, $start, $out, |v| v, $step, |b: &[$t], t: $t| b.iter().fold(t, |m, &v| ($step)(m, v)) == t) };
734 (@grouped $x:ident, $start:ident, $out:ident, $t:ty, $step:expr) => { grouped($x, $start, $out, |v| v, $step, |t| t, true) };
735 (signed $($t:ty),+) => {$( closed!(@one Xnor, $t, !0, |a: $t, b: $t| !(a ^ b), grouped, |x: &[$t]| xnor_of(x.iter().fold(0, |a: $t, &b| a ^ b), x.len())); )+};
736 ($($t:ty),+) => {$(
737 closed!(@one Max, $t, <$t>::MIN, Ord::max, records);
738 closed!(@one Min, $t, <$t>::MAX, Ord::min, records);
739 closed!(@one And, $t, !0, |a: $t, b: $t| a & b, grouped);
740 closed!(@one Or, $t, 0, |a: $t, b: $t| a | b, grouped);
741 closed!(@one Xor, $t, 0, |a: $t, b: $t| a ^ b, grouped);
742 )+};
743}
744closed!(i16, i32, i64, u8, u16, u32, u64);
745closed!(signed i16, i32, i64);
746impl Reduction<u8> for Xnor {
748 type Total = i16;
749 const COMMUTATIVE: bool = true;
750 #[inline(always)]
751 fn unit() -> i16 { !0 }
752 #[inline(always)]
753 fn step(acc: i16, x: u8) -> i16 { !(acc ^ i16::from(x)) }
754 #[inline(always)]
755 fn merge(a: i16, b: i16) -> (i16, bool) { (!(a ^ b), true) }
756 #[inline]
757 fn slice(x: &[u8]) -> (i16, bool) { (xnor_of(x.iter().fold(0i16, |a, &b| a ^ i16::from(b)), x.len()), true) }
758 #[inline]
759 fn scan(x: &[u8], start: Option<i16>, out: &mut Vec<i16>) -> bool { grouped(x, start, out, i16::from, |a: i16, b: i16| !(a ^ b), |t| t, true); true }
760 #[inline]
761 fn rows(data: &[u8], columns: usize, out: &mut [i16]) -> bool { rows_side_by_side::<u8, Self>(data, columns, out) }
762}
763macro_rules! booleans {
765 ($($($r:ident)|+: $unit:expr, |$a:ident, $b:ident| $step:expr, $table:expr;)+) => {$($(
766 impl Reduction<bool> for $r {
767 type Total = bool;
768 const COMMUTATIVE: bool = true;
769 #[inline(always)]
770 fn unit() -> bool { $unit }
771 #[inline(always)]
772 fn step($a: bool, $b: bool) -> bool { $step }
773 #[inline(always)]
774 fn merge(a: bool, b: bool) -> (bool, bool) { (<Self as Reduction<bool>>::step(a, b), true) }
775 #[inline]
776 fn slice(x: &[bool]) -> (bool, bool) { (folded(x, $unit, <Self as Reduction<bool>>::step, <Self as Reduction<bool>>::step), true) }
777 #[inline]
778 fn scan(x: &[bool], start: Option<bool>, out: &mut Vec<bool>) -> bool {
779 Logical($table).scan(x, start, out);
780 true
781 }
782 #[inline]
783 fn short_lanes(data: &[bool], len: usize, start: Option<bool>, out: &mut Vec<bool>) -> Option<bool> {
784 (1..=64).contains(&len).then(|| { Logical($table).short_lanes(data, len, start, out); true })
785 }
786 }
787 )+)+};
788}
789booleans! {
790 Max | Or: false, |a, b| a | b, [[false, true], [true, true]];
791 Min | And: true, |a, b| a & b, [[false, false], [false, true]];
792 Xor: false, |a, b| a ^ b, [[false, true], [true, false]];
793 Xnor: true, |a, b| a == b, [[true, false], [false, true]];
794}
795
796pub(crate) fn records<T: Copy + PartialEq>(
798 x: &[T], start: Option<T>, out: &mut Vec<T>, from: impl Fn(T) -> T + Copy, step: impl Fn(T, T) -> T + Copy, keeps: impl Fn(&[T], T) -> bool) {
799 let Some(&first) = x.first() else { return };
800 let first = from(first);
801 let (mut total, mut at, mut changes) = (start.map_or(first, |s| step(s, first)), 0, 0);
802 while at < x.len() {
803 let test = changes < 4;
804 let block = &x[at..x.len().min(at + if test { 256 } else { 256 << (changes - 4).min(6) })];
805 at += block.len();
806 if test && keeps(block, total) {
807 out.resize(out.len() + block.len(), total);
808 changes = 0;
809 } else {
810 grouped(block, Some(total), out, from, step, |t| t, false);
811 let last = out[out.len() - 1];
812 (total, changes) = (last, if last == total { 0 } else { changes + 1 });
813 }
814 }
815}
816
817#[inline(always)]
819pub(crate) fn below<F: PartialOrd + Copy>(block: &[F], total: F) -> bool { block.iter().fold(true, |k, &v| k & (v < total)) }
820
821#[inline(always)]
823pub(crate) fn above<F: PartialOrd + Copy>(block: &[F], total: F) -> bool { block.iter().fold(true, |k, &v| k & (v > total)) }
824
825const PIECE: usize = 1 << 18;
827const ROWS: usize = 16;
829
830fn piece_size(inner: usize) -> usize { PIECE.max(ROWS * inner).next_multiple_of(CHUNK) }
832
833#[inline(always)]
835fn one_of<T, R: Reduction<T>>(x: T, ok: &mut bool) -> R::Total {
836 let (t, fits) = R::one(x);
837 *ok &= fits;
838 t
839}
840
841#[inline(always)]
843fn steps_of<T: Copy, R: Reduction<T>>(totals: &mut [R::Total], x: &[T]) -> bool {
844 let (n, mut fits) = (totals.len().min(x.len()), true);
845 for k in 0..n { let (t, f) = R::checked(totals[k], x[k]); (totals[k], fits) = (t, fits & f) }
846 fits
847}
848
849#[inline(always)]
851fn started<T: Copy, R: Reduction<T>>(out: &mut Vec<R::Total>, rows: &[&[T]]) -> bool {
852 #[inline(always)]
853 fn step<T, R: Reduction<T>>((t, f): (R::Total, bool), x: T) -> (R::Total, bool) { let (u, g) = R::checked(t, x); (u, f & g) }
854 let mut fits = true;
855 match *rows {
856 [] => {}
857 [a] => out.extend(a.iter().map(|&x| {
858 let (t, f) = R::one(x);
859 fits &= f;
860 t
861 })),
862 [a, b] => out.extend(a.iter().zip(b).map(|(&x, &y)| {
863 let (t, f) = step::<T, R>(R::one(x), y);
864 fits &= f;
865 t
866 })),
867 [a, b, c, ..] => out.extend(a.iter().zip(b).zip(c).map(|((&x, &y), &z)| {
868 let (t, f) = step::<T, R>(step::<T, R>(R::one(x), y), z);
869 fits &= f;
870 t
871 })),
872 }
873 fits
874}
875
876#[inline(always)]
878fn rows_folded<'a, T: Copy + 'a, R: Reduction<T>>(out: &mut Vec<R::Total>, width: usize, row: impl Fn(usize) -> &'a [T]) -> bool {
879 let at = out.len();
880 let mut fits = match width {
881 1 => started::<T, R>(out, &[row(0)]), 2 => started::<T, R>(out, &[row(0), row(1)]), _ => started::<T, R>(out, &[row(0), row(1), row(2)])
882 };
883 for i in 3..width { fits &= steps_of::<T, R>(&mut out[at..], row(i)); }
884 fits
885}
886
887#[inline(always)]
889fn merges_of<T, R: Reduction<T>>(totals: &mut [R::Total], parts: &[R::Total]) -> bool {
890 let (n, mut fits) = (totals.len().min(parts.len()), true);
891 for k in 0..n { let (t, f) = R::merge(totals[k], parts[k]); (totals[k], fits) = (t, fits & f) }
892 fits
893}
894
895#[inline(always)]
897fn merge_of<T, R: Reduction<T>>(a: R::Total, b: R::Total, ok: &mut bool) -> R::Total {
898 let (t, fits) = R::merge(a, b);
899 *ok &= fits;
900 t
901}
902
903#[inline(always)]
905fn chain_of<T, R: Reduction<T>>(a: R::Total, b: R::Total, ok: &mut bool) -> R::Total {
906 let (t, fits) = R::chain(a, b);
907 *ok &= fits;
908 t
909}
910
911const SHORT_LANE: usize = 64;
913
914const LANES: usize = 8;
916
917#[inline(always)]
919fn side_by_side<const B: usize, T: Copy, S: Copy>(
920 data: &[T], len: usize, first: impl Fn(T) -> (S, bool), step: impl Fn(S, T) -> (S, bool), mut each: impl FnMut(usize, S), mut last: impl FnMut(usize, S)
921) -> bool {
922 let mut fits = true;
923 for (g, group) in data.chunks_exact(B * len).enumerate() {
924 let at = g * B * len;
925 let mut ok = [true; B];
926 let mut t: [S; B] = std::array::from_fn(|b| {
927 let (t, f) = first(group[b * len]);
928 ok[b] = f;
929 each(at + b * len, t);
930 t
931 });
932 for i in 1..len {
933 for b in 0..B {
934 let (s, f) = step(t[b], unsafe { *group.get_unchecked(b * len + i) });
936 (t[b], ok[b]) = (s, ok[b] & f);
937 each(at + b * len + i, s);
938 }
939 }
940 for (b, &t) in t.iter().enumerate() { last(g * B + b, t) }
941 fits &= ok.into_iter().all(|f| f);
942 }
943 fits
944}
945
946pub trait Term: Copy {
948 fn term<C: Tally>(self) -> C;
950 #[inline]
952 fn terms<C: Tally>(items: &[Self]) -> C { items.iter().fold(C::default(), |s, &x| s.add(x.term())) }
953}
954impl Term for bool {
955 #[inline(always)]
956 fn term<C: Tally>(self) -> C { C::low(self as i64) }
957 #[inline]
959 fn terms<C: Tally>(items: &[bool]) -> C { items.chunks(192).fold(C::default(), |s, c| s.add(C::low(c.iter().fold(0u8, |n, &b| n + u8::from(b)).into()))) }
960}
961macro_rules! terms {
962 ($($t:ty),+) => {$(
963 impl Term for $t {
964 #[inline(always)]
965 fn term<C: Tally>(self) -> C { C::low(self as i64) }
966 }
967 )+};
968}
969terms!(i16, i32, i64, u8);
970
971pub trait Tally: Copy + Default + Send + Sync {
973 fn add(self, y: Self) -> Self;
975 fn low(n: i64) -> Self;
977}
978macro_rules! tallies {
979 ($($t:ty),+) => {$(
980 impl Tally for $t {
981 #[inline(always)]
982 fn add(self, y: $t) -> $t { self.wrapping_add(y) }
983 #[inline(always)]
984 fn low(n: i64) -> $t { n as $t }
985 }
986 )+};
987}
988tallies!(u8, i16, i32, i64);
989
990#[derive(Clone, Copy, Debug)]
992pub struct SumAs<C>(std::marker::PhantomData<C>);
993impl<T: Term, C: Tally> Reduction<T> for SumAs<C> {
994 type Total = C;
995 const COMMUTATIVE: bool = true;
996 #[inline(always)]
997 fn unit() -> C { C::default() }
998 #[inline(always)]
999 fn step(acc: C, x: T) -> C { acc.add(x.term()) }
1000 #[inline(always)]
1001 fn merge(a: C, b: C) -> (C, bool) { (a.add(b), true) }
1002 #[inline]
1003 fn slice(x: &[T]) -> (C, bool) { (T::terms(x), true) }
1004 #[inline]
1005 fn scan(x: &[T], start: Option<C>, out: &mut Vec<C>) -> bool { grouped(x, start, out, T::term, C::add, |t| t, size_of::<C>() < 8); true }
1006}
1007
1008#[derive(Clone, Copy, Debug)]
1010pub struct Squares;
1011
1012#[inline(always)]
1014fn add_square<T: Scalar>(s: f64, x: T) -> f64 {
1015 let re = x.re().wide();
1016 let s = re.mul_acc(re, s);
1017 if T::COMPLEX { let im = x.im().wide(); im.mul_acc(im, s) } else { s }
1018}
1019
1020impl<T: Scalar> Reduction<T> for Squares {
1021 type Total = f64;
1022 const COMMUTATIVE: bool = true;
1023 #[inline(always)]
1024 fn unit() -> f64 { 0.0 }
1025 #[inline(always)]
1026 fn step(acc: f64, x: T) -> f64 { add_square(acc, x) }
1027 #[inline(always)]
1028 fn merge(a: f64, b: f64) -> (f64, bool) { (a + b, true) }
1029 #[inline]
1030 fn slice(x: &[T]) -> (f64, bool) { (folded(x, 0.0, add_square, |a, b| a + b), true) }
1031 #[inline]
1032 fn scan(x: &[T], start: Option<f64>, out: &mut Vec<f64>) -> bool { grouped(x, start, out, |v| add_square(0.0, v), |a, b| a + b, |t| t, false); true }
1033}
1034
1035#[inline(always)]
1037fn pairwise<A: Copy, S: Copy>(x: &[A], unit: S, pair: impl Fn(S, A, A) -> S, merge: impl Fn(S, S) -> S, last: impl Fn(S, A) -> S) -> S {
1038 let (pairs, rest) = x.as_chunks::<2>();
1039 let total = folded(pairs, unit, |s, [a, b]| pair(s, a, b), merge);
1040 rest.first().map_or(total, |&z| last(total, z))
1041}
1042
1043#[inline(always)]
1045fn quotient<T: Copy, R: Native>(
1046 x: &[T], one: T, divide: impl Fn(T, T) -> T, times: impl Fn(T, T) -> T + Copy, size: impl Fn(T) -> R + Copy, finite: impl Fn(T) -> bool) -> (T, bool) {
1047 let step = move |(p, low): (T, R), q: T| { let r = times(p, q); (r, low.min(size(r))) };
1048 let (total, low) = pairwise(x, (one, R::one()), |s, a, b| step(s, divide(a, b)), move |s, (q, l)| {
1049 let (r, m) = step(s, q);
1050 (r, m.min(l))
1051 }, step);
1052 (total, low >= R::min_positive_value() && finite(total))
1053}
1054
1055#[inline(always)]
1057fn alternated<A: Copy, S: Copy + std::ops::Sub<Output = S>>(x: &[A], unit: S, add: impl Fn(S, A) -> S + Copy, merge: impl Fn(S, S) -> S) -> S {
1058 let merge = &merge;
1059 running_results(x, unit, add, merge, move |parts, rest| {
1060 let ((pairs, _), (tail, last)) = (parts.as_chunks::<2>(), rest.as_chunks::<2>());
1061 let (even, odd) = pairs.iter().fold((unit, unit), |(e, o), &[p, q]| (merge(e, p), merge(o, q)));
1062 let (even, odd) = tail.iter().fold((even, odd), |(e, o), &[p, q]| (add(e, p), add(o, q)));
1063 last.first().map_or(even, |&z| add(even, z)) - odd
1064 })
1065}
1066
1067fn checked_steps<T: Copy, R: Reduction<T>>(x: &[T], start: Option<R::Total>, out: &mut Vec<R::Total>) -> bool {
1069 let (mut total, mut fits) = (start.unwrap_or_else(R::identity), true);
1070 out.extend(x.iter().map(|&v| {
1071 let ok;
1072 (total, ok) = R::checked(total, v);
1073 fits &= ok;
1074 total
1075 }));
1076 fits
1077}
1078
1079#[derive(Clone, Copy, Debug, Default)]
1081pub struct Alternated<R>(std::marker::PhantomData<R>);
1082
1083pub trait Inverse<T>: Reduction<T> {
1085 fn inverse(acc: Self::Total, x: T) -> (Self::Total, bool);
1087 fn undo(a: Self::Total, b: Self::Total) -> (Self::Total, bool);
1089 #[inline(always)]
1091 fn fits(_: Self::Total) -> bool { true }
1092 fn alternating(x: &[T]) -> (Self::Total, bool);
1094}
1095
1096impl<T: Copy, R: Inverse<T>> Reduction<T> for Alternated<R> {
1097 type Total = (R::Total, bool);
1098 #[inline(always)]
1099 fn unit() -> Self::Total { (R::unit(), false) }
1100 #[inline(always)]
1101 fn identity() -> Self::Total { (R::identity(), false) }
1102 #[inline(always)]
1103 fn step((s, odd): Self::Total, x: T) -> Self::Total { (if odd { R::inverse(s, x).0 } else { R::step(s, x) }, !odd) }
1104 #[inline(always)]
1105 fn checked((s, odd): Self::Total, x: T) -> (Self::Total, bool) {
1106 let (t, ok) = if odd { R::inverse(s, x) } else { R::checked(s, x) };
1107 ((t, !odd), ok && R::fits(t))
1108 }
1109 #[inline(always)]
1110 fn merge((a, odd): Self::Total, (b, more): Self::Total) -> (Self::Total, bool) {
1111 let (t, ok) = if odd { R::undo(a, b) } else { R::merge(a, b) };
1112 ((t, odd ^ more), ok && R::fits(t))
1113 }
1114 #[inline]
1115 fn slice(x: &[T]) -> (Self::Total, bool) { let (t, ok) = R::alternating(x); ((t, x.len() % 2 == 1), ok) }
1116}
1117
1118impl<F: Native> Inverse<F> for Sum {
1119 #[inline(always)]
1120 fn inverse(acc: F, x: F) -> (F, bool) { (acc.algebraic_add(-x), true) }
1121 #[inline(always)]
1122 fn undo(a: F, b: F) -> (F, bool) { (a.algebraic_add(-b), true) }
1123 #[inline]
1124 fn alternating(x: &[F]) -> (F, bool) { (alternated(x, F::zero(), |s, v| s + v, |a, b| a + b), true) }
1125}
1126
1127impl<F: Native> Inverse<Complex<F>> for Sum {
1128 #[inline(always)]
1129 fn inverse(acc: Complex<F>, x: Complex<F>) -> (Complex<F>, bool) { (<Sum as Reduction<Complex<F>>>::step(acc, -x), true) }
1130 #[inline(always)]
1131 fn undo(a: Complex<F>, b: Complex<F>) -> (Complex<F>, bool) { (<Sum as Reduction<Complex<F>>>::step(a, -b), true) }
1132 #[inline]
1133 fn alternating(x: &[Complex<F>]) -> (Complex<F>, bool) { (alternated(x, Complex::new(F::zero(), F::zero()), |s, v| s + v, |a, b| a + b), true) }
1134}
1135
1136macro_rules! alternated_integers {
1138 ($($t:ty => $acc:ty, $run:expr);+) => {$(
1139 impl Inverse<$t> for Sum {
1140 #[inline(always)]
1141 fn inverse(acc: i64, x: $t) -> (i64, bool) { let (t, over) = acc.overflowing_sub(i64::from(x)); (t, !over) }
1142 #[inline(always)]
1143 fn undo(a: i64, b: i64) -> (i64, bool) { let (t, over) = a.overflowing_sub(b); (t, !over) }
1144 #[inline]
1145 fn alternating(x: &[$t]) -> (i64, bool) {
1146 let run = |c: &[$t]| i128::from(alternated(c, 0 as $acc, |s, v| s + <$acc>::from(v), |s, t| s + t));
1147 let total: i128 = x.chunks($run).map(run).sum();
1148 (total as i64, i64::try_from(total).is_ok())
1149 }
1150 }
1151 )+};
1152}
1153alternated_integers!(i16 => i32, 1 << 15; u8 => i32, 1 << 15; i32 => i64, 1 << 31; i64 => i128, 1 << 31);
1154
1155#[inline(always)]
1157fn normal<F: Native>(q: F) -> bool { q.abs() >= F::min_positive_value() && q.is_finite() }
1158
1159#[derive(Clone, Copy, Debug)]
1161pub struct Scaled<F> { m: F, e: F, normal: bool }
1162
1163impl<F: Native> Scaled<F> {
1164 #[inline(always)]
1166 fn of(p: F, e: F, normal: bool) -> Self {
1167 let (lo, hi) = if F::WIDE { (1.4916681462400413e-154, 6.703903964971299e153) } else { (1.0842021724855044e-19, 9.223372036854776e18) };
1168 if p.abs() >= F::narrow(lo) && p.abs() <= F::narrow(hi) { Self { m: p, e, normal } } else { Self::rescaled(p, e, normal) }
1169 }
1170 #[inline]
1172 fn rescaled(p: F, e: F, normal: bool) -> Self {
1173 let (m, f) = p.abs().split();
1174 Self { m: if p < F::zero() { -m } else { m }, e: e + f, normal: normal & p.abs().positive_normal() }
1175 }
1176 #[inline(always)]
1178 fn times(self, q: F) -> Self { Self::of(self.m * q, self.e, self.normal) }
1179 #[inline(always)]
1181 fn over(self, q: F) -> Self { Self::of(self.m / q, self.e, self.normal) }
1182 #[inline(always)]
1184 fn joined(self, b: Self) -> Self { Self::of(self.m * b.m, self.e + b.e, self.normal & b.normal) }
1185 #[inline(always)]
1187 fn undone(self, b: Self) -> Self { Self::of(self.m / b.m, self.e - b.e, self.normal & b.normal) }
1188 #[inline]
1190 pub fn value(self) -> F {
1191 let Self { m, e, .. } = Self::rescaled(self.m, self.e, true);
1192 let far = F::narrow(if F::WIDE { 2040.0 } else { 250.0 });
1193 crate::float::pow2_times::<F, { crate::FUSED }>(m, e.max(-far).min(far))
1194 }
1195}
1196
1197#[derive(Clone, Copy, Debug, Default)]
1199pub struct ScaledProduct;
1200
1201impl<F: Native> Reduction<F> for ScaledProduct {
1202 type Total = Scaled<F>;
1203 const COMMUTATIVE: bool = true;
1204 #[inline(always)]
1205 fn unit() -> Scaled<F> { Scaled { m: F::one(), e: F::zero(), normal: true } }
1206 #[inline(always)]
1207 fn step(acc: Scaled<F>, x: F) -> Scaled<F> { acc.times(x) }
1208 #[inline(always)]
1209 fn checked(acc: Scaled<F>, x: F) -> (Scaled<F>, bool) { let t = acc.times(x); (t, t.normal) }
1210 #[inline(always)]
1211 fn merge(a: Scaled<F>, b: Scaled<F>) -> (Scaled<F>, bool) { let t = a.joined(b); (t, t.normal) }
1212}
1213
1214impl<F: Native> Inverse<F> for ScaledProduct {
1215 #[inline(always)]
1216 fn inverse(acc: Scaled<F>, x: F) -> (Scaled<F>, bool) { let t = acc.over(x); (t, t.normal) }
1217 #[inline(always)]
1218 fn undo(a: Scaled<F>, b: Scaled<F>) -> (Scaled<F>, bool) { let t = a.undone(b); (t, t.normal) }
1219 #[inline]
1220 fn alternating(x: &[F]) -> (Scaled<F>, bool) {
1221 let (q, fits) = quotient(x, F::one(), |a, b| a / b, F::algebraic_mul, F::abs, F::is_finite);
1223 let t = if fits { Scaled::of(q, F::zero(), true) }
1224 else { pairwise(x, <Self as Reduction<F>>::unit(), |s, a, b| s.times(a / b), Scaled::joined, Scaled::times) };
1225 (t, t.normal)
1226 }
1227}
1228
1229#[inline(always)]
1231fn complex_times<F: Native>(p: Complex<F>, q: Complex<F>) -> Complex<F> {
1232 Complex::new(p.re.algebraic_mul(q.re).algebraic_add(-p.im.algebraic_mul(q.im)), p.re.algebraic_mul(q.im).algebraic_add(p.im.algebraic_mul(q.re)))
1233}
1234
1235#[inline(always)]
1237fn complex_size<F: Native>(p: Complex<F>) -> F { p.re.abs().max(p.im.abs()) }
1238
1239impl<F: Native> Inverse<Complex<F>> for Product {
1241 #[inline(always)]
1242 fn inverse(acc: Complex<F>, x: Complex<F>) -> (Complex<F>, bool) { (crate::divide(acc, x), true) }
1243 #[inline(always)]
1244 fn undo(a: Complex<F>, b: Complex<F>) -> (Complex<F>, bool) { (crate::divide(a, b), true) }
1245 #[inline(always)]
1246 fn fits(q: Complex<F>) -> bool { normal(complex_size(q)) }
1247 #[inline]
1248 fn alternating(x: &[Complex<F>]) -> (Complex<F>, bool) {
1249 quotient(x, Complex::new(F::one(), F::zero()), crate::divide, complex_times, complex_size, |p| p.re.is_finite() & p.im.is_finite())
1250 }
1251}
1252
1253#[macro_export]
1255macro_rules! with_reduce {
1256 ($r:expr, |$t:ident| $body:expr) => {
1257 match $r {
1258 $crate::reduce::Reduce::Sum => { type $t = $crate::reduce::Sum; $body }
1259 $crate::reduce::Reduce::Product => { type $t = $crate::reduce::Product; $body }
1260 $crate::reduce::Reduce::Max => { type $t = $crate::reduce::Max; $body }
1261 $crate::reduce::Reduce::Min => { type $t = $crate::reduce::Min; $body }
1262 $crate::reduce::Reduce::And => { type $t = $crate::reduce::And; $body }
1263 $crate::reduce::Reduce::Or => { type $t = $crate::reduce::Or; $body }
1264 $crate::reduce::Reduce::Xor => { type $t = $crate::reduce::Xor; $body }
1265 $crate::reduce::Reduce::Xnor => { type $t = $crate::reduce::Xnor; $body }
1266 }
1267 };
1268}
1269
1270#[doc(hidden)]
1272#[macro_export]
1273macro_rules! native_reduction {
1274 (@inverse $kind:expr, $dtype:expr, |$R:ident, $T:ident| $body:expr) => {{
1275 use $crate::buffer::{DType, FloatWidth};
1276 use $crate::reduce::{Alternated, Product, Reduce, ScaledProduct, Sum};
1277 match ($dtype, $kind) {
1278 (DType::Integer(w), Reduce::Sum) => with_width!(w, I => $crate::native_reduction!(@at Alternated<Sum>, I, |$R, $T| $body)),
1279 (DType::Float(FloatWidth::F16 | FloatWidth::F32), Reduce::Sum) => $crate::native_reduction!(@at Alternated<Sum>, f32, |$R, $T| $body),
1280 (DType::Float(FloatWidth::F64), Reduce::Sum) => $crate::native_reduction!(@at Alternated<Sum>, f64, |$R, $T| $body),
1281 (DType::Float(FloatWidth::F16 | FloatWidth::F32), Reduce::Product) => $crate::native_reduction!(@at Alternated<ScaledProduct>, f32, |$R, $T| $body),
1282 (DType::Float(FloatWidth::F64), Reduce::Product) => $crate::native_reduction!(@at Alternated<ScaledProduct>, f64, |$R, $T| $body),
1283 (DType::Complex(w), Reduce::Sum) => with_complex_width!(w, C => $crate::native_reduction!(@at Alternated<Sum>, C, |$R, $T| $body)),
1284 (DType::Complex(w), Reduce::Product) => with_complex_width!(w, C => $crate::native_reduction!(@at Alternated<Product>, C, |$R, $T| $body)),
1285 _ => None,
1286 }
1287 }};
1288 (@at $r:ty, $t:ty, |$R:ident, $T:ident| $body:expr) => {{ type $R = $r; type $T = $t; $body }};
1289 ($kind:expr, $dtype:expr, $checked:expr, |$R:ident, $T:ident| $body:expr) => {{
1290 use $crate::buffer::DType;
1291 use $crate::reduce::Reduce;
1292 match ($dtype, $kind) {
1293 (DType::Boolean, Reduce::Sum) => $crate::native_reduction!(@at $crate::reduce::SumAs<i64>, bool, |$R, $T| $body),
1294 (DType::Boolean, Reduce::Product | Reduce::And | Reduce::Min) => $crate::native_reduction!(@at $crate::reduce::And, bool, |$R, $T| $body),
1295 (DType::Boolean, Reduce::Or | Reduce::Max) => $crate::native_reduction!(@at $crate::reduce::Or, bool, |$R, $T| $body),
1296 (DType::Boolean, Reduce::Xor) => $crate::native_reduction!(@at $crate::reduce::Xor, bool, |$R, $T| $body),
1297 (DType::Boolean, Reduce::Xnor) => $crate::native_reduction!(@at $crate::reduce::Xnor, bool, |$R, $T| $body),
1298 (DType::Integer(w), Reduce::Sum) if !$checked => with_width!(w, I => $crate::native_reduction!(@at $crate::Wrapped<$crate::reduce::Sum>, I, |$R, $T| $body)),
1299 (DType::Integer(w), Reduce::Sum) => with_width!(w, I => $crate::native_reduction!(@at $crate::reduce::Sum, I, |$R, $T| $body)),
1300 (DType::Integer(w), Reduce::Product) if !$checked => with_width!(w, I => $crate::native_reduction!(@at $crate::Wrapped<$crate::reduce::Product>, I, |$R, $T| $body)),
1301 (DType::Integer(w), Reduce::Product) => with_width!(w, I => $crate::native_reduction!(@at $crate::reduce::Product, I, |$R, $T| $body)),
1302 (DType::Integer(w), Reduce::Max) => with_width!(w, I => $crate::native_reduction!(@at $crate::reduce::Max, I, |$R, $T| $body)),
1303 (DType::Integer(w), Reduce::Min) => with_width!(w, I => $crate::native_reduction!(@at $crate::reduce::Min, I, |$R, $T| $body)),
1304 (DType::Integer(w), Reduce::And) => with_width!(w, I => $crate::native_reduction!(@at $crate::reduce::And, I, |$R, $T| $body)),
1305 (DType::Integer(w), Reduce::Or) => with_width!(w, I => $crate::native_reduction!(@at $crate::reduce::Or, I, |$R, $T| $body)),
1306 (DType::Integer(w), Reduce::Xor) => with_width!(w, I => $crate::native_reduction!(@at $crate::reduce::Xor, I, |$R, $T| $body)),
1307 (DType::Integer(w), Reduce::Xnor) => with_width!(w, I => $crate::native_reduction!(@at $crate::reduce::Xnor, I, |$R, $T| $body)),
1308 (DType::Float(w), Reduce::Sum) => with_float_width!(w, F => $crate::native_reduction!(@at $crate::reduce::Sum, F, |$R, $T| $body)),
1309 (DType::Float(w), Reduce::Product) => with_float_width!(w, F => $crate::native_reduction!(@at $crate::reduce::Product, F, |$R, $T| $body)),
1310 (DType::Float(w), Reduce::Max) => with_float_width!(w, F => $crate::native_reduction!(@at $crate::reduce::Max, F, |$R, $T| $body)),
1311 (DType::Float(w), Reduce::Min) => with_float_width!(w, F => $crate::native_reduction!(@at $crate::reduce::Min, F, |$R, $T| $body)),
1312 (DType::Complex(w), Reduce::Sum) => with_complex_width!(w, C => $crate::native_reduction!(@at $crate::reduce::Sum, C, |$R, $T| $body)),
1313 (DType::Complex(w), Reduce::Product) => with_complex_width!(w, C => $crate::native_reduction!(@at $crate::reduce::Product, C, |$R, $T| $body)),
1314 _ => None,
1315 }
1316 }};
1317}
1318pub use crate::native_reduction;
1319
1320macro_rules! wrapped_products {
1321 ($($t:ty),+) => {$(
1322 impl Reduction<$t> for crate::Wrapped<Product> {
1323 type Total = i64;
1324 const COMMUTATIVE: bool = true;
1325 fn unit() -> i64 { 1 }
1326 fn step(total: i64, x: $t) -> i64 { total.wrapping_mul(x.into()) }
1327 fn merge(x: i64, y: i64) -> (i64, bool) { (x.wrapping_mul(y), true) }
1328 fn slice(x: &[$t]) -> (i64, bool) { (folded(x, 1i64, |total, x| total.wrapping_mul(x.into()), i64::wrapping_mul), true) }
1329 fn scan(x: &[$t], start: Option<i64>, out: &mut Vec<i64>) -> bool {
1330 grouped(x, start, out, |x| x.into(), i64::wrapping_mul, |x| x, true);
1331 true
1332 }
1333 }
1334 )+};
1335}
1336wrapped_products!(u8, i16, i32, i64);
1337
1338mod along;
1339mod pieces;
1340mod run;
1341mod scans;
1342mod sliding;
1343pub use {along::*, pieces::*, run::*, scans::*, sliding::*};