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 = 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 store(o, add(c, p0));
537 store(o.add(4), add(add(c, s0), p1));
538 store(o.add(8), add(add(c, s01), p2));
539 store(o.add(12), add(add(c, s012), p3));
540 c = add(c, add(s012, last(p3)));
541 }
542 let mut total = first(c);
543 for (o, &v) in tail.iter_mut().zip(rest) {
544 total += v;
545 o.write(total);
546 }
547 })
548 }
549 }
550 #[cfg(target_arch = "wasm32")]
551 lanes! {
552 f64: 2, f64x2_splat, v128_load, f64x2_pmax, f64x2_pmin, lane(f64x2_extract_lane), fold, grouped;
553 f32: 4, f32x4_splat, v128_load, f32x4_pmax, f32x4_pmin, lane(f32x4_extract_lane), fold, vector;
554 }
555}
556
557#[cfg(not(any(target_arch = "aarch64", all(target_arch = "wasm32", target_feature = "simd128"))))]
559mod plain {
560 #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
561 use core::arch::x86_64::*;
562 macro_rules! plain {
563 ($($t:ty: $lanes:literal, $load:ident, $min:ident, $max:ident;)+) => {$(
564 impl super::Folds for $t {
565 #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
566 #[inline(always)]
567 fn ordered_lanes<const W: usize>(a: [$t; W], c: [$t; W]) -> ([$t; W], [$t; W]) {
568 const { assert!(W == $lanes) };
569 unsafe {
571 let (x, y) = ($load(a.as_ptr()), $load(c.as_ptr()));
572 (core::mem::transmute_copy(&$min(y, x)), core::mem::transmute_copy(&$max(x, y)))
573 }
574 }
575 }
576 )+};
577 }
578 plain! {
579 f32: 8, _mm256_loadu_ps, _mm256_min_ps, _mm256_max_ps;
580 f64: 4, _mm256_loadu_pd, _mm256_min_pd, _mm256_max_pd;
581 }
582}
583
584#[inline(always)]
586fn narrow_sum<T: Copy + Into<i64>>(x: &[T]) -> (i64, bool) {
587 let total: i128 = x.chunks(1 << 31).map(|c| i128::from(folded(c, 0i64, |s, v| s + v.into(), |a, b| a + b))).sum();
588 (total as i64, i64::try_from(total).is_ok())
589}
590
591#[inline(always)]
593fn small_sum<T: Copy + Into<i32>>(x: &[T]) -> (i64, bool) {
594 (x.chunks(1 << 15).map(|c| i64::from(folded(c, 0i32, |s, v| s + v.into(), |a, b| a + b))).sum(), true)
595}
596
597#[inline(always)]
599fn wide_sum(x: &[i64]) -> (i64, bool) {
600 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));
601 let total: i128 = x.chunks(1 << 31).map(|c| {
602 let (h, l) = halves(c);
603 (i128::from(h) << 32) + i128::from(l)
604 }).sum();
605 (total as i64, i64::try_from(total).is_ok())
606}
607
608#[inline(always)]
610fn signed_cumsum<T: Copy + Into<i64>>(x: &[T], start: i64, out: &mut Vec<i64>) -> bool {
611 unsafe {
612 appended(out, x.len(), |slots| {
613 let (mut total, mut overflow) = (start, 0i64);
614 for (o, &v) in slots.iter_mut().zip(x) {
615 let v: i64 = v.into();
616 let sum = total.wrapping_add(v);
617 overflow |= (total ^ sum) & (v ^ sum);
618 total = sum;
619 o.write(sum);
620 }
621 overflow >= 0
622 })
623 }
624}
625
626#[inline(always)]
628fn checked_scan<T: Copy + Into<S>, S: Copy>(x: &[T], start: S, out: &mut Vec<S>, op: impl Fn(S, S) -> (S, bool)) -> bool {
629 unsafe {
630 appended(out, x.len(), |slots| {
631 let (mut total, mut ok) = (start, true);
632 for (o, &v) in slots.iter_mut().zip(x) {
633 let (next, fits) = op(total, v.into());
634 (total, ok) = (next, ok & fits);
635 o.write(next);
636 }
637 ok
638 })
639 }
640}
641
642macro_rules! integers {
644 ($($t:ty: $sum:ident),+) => {$(
645 impl Reduction<$t> for Sum {
646 type Total = i64;
647 const COMMUTATIVE: bool = true;
648 #[inline(always)]
649 fn unit() -> i64 { 0 }
650 #[inline(always)]
651 fn step(acc: i64, x: $t) -> i64 { acc.plus(x.into()).0 }
652 #[inline(always)]
653 fn checked(acc: i64, x: $t) -> (i64, bool) { acc.plus(x.into()) }
654 #[inline(always)]
655 fn merge(a: i64, b: i64) -> (i64, bool) { a.plus(b) }
656 #[inline]
657 fn slice(x: &[$t]) -> (i64, bool) { $sum(x) }
658 #[inline]
659 fn scan(x: &[$t], start: Option<i64>, out: &mut Vec<i64>) -> bool { signed_cumsum(x, start.unwrap_or(0), out) }
660 }
661 impl Reduction<$t> for Product {
662 type Total = i64;
663 const COMMUTATIVE: bool = true;
664 #[inline(always)]
665 fn unit() -> i64 { 1 }
666 #[inline(always)]
667 fn step(acc: i64, x: $t) -> i64 { acc.times(x.into()).0 }
668 #[inline(always)]
669 fn checked(acc: i64, x: $t) -> (i64, bool) { acc.times(x.into()) }
670 #[inline(always)]
671 fn merge(a: i64, b: i64) -> (i64, bool) { a.times(b) }
672 #[inline]
673 fn slice(x: &[$t]) -> (i64, bool) {
674 x.iter().fold((1 as i64, true), |(p, ok), &v| { let (q, fits) = p.times(v.into()); (q, ok & fits) })
675 }
676 #[inline]
677 fn scan(x: &[$t], start: Option<i64>, out: &mut Vec<i64>) -> bool { checked_scan(x, start.unwrap_or(1), out, <i64>::times) }
678 }
679 )+};
680}
681integers!(u8: small_sum, i16: small_sum, i32: narrow_sum, i64: wide_sum);
682
683macro_rules! wrapped_sums {
685 ($($t:ty: $slice:expr),+) => {$(
686 impl Reduction<$t> for crate::Wrapped<Sum> {
687 type Total = i64;
688 const COMMUTATIVE: bool = true;
689 #[inline(always)]
690 fn unit() -> i64 { 0 }
691 #[inline(always)]
692 fn step(acc: i64, x: $t) -> i64 { acc.wrapping_add(x.into()) }
693 #[inline(always)]
694 fn merge(a: i64, b: i64) -> (i64, bool) { (a.wrapping_add(b), true) }
695 #[inline]
696 fn slice(x: &[$t]) -> (i64, bool) { (($slice)(x), true) }
697 #[inline]
698 fn scan(x: &[$t], start: Option<i64>, out: &mut Vec<i64>) -> bool { signed_cumsum(x, start.unwrap_or(0), out); true }
699 }
700 )+};
701}
702wrapped_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));
703
704fn xnor_of<T: std::ops::Not<Output = T>>(xor: T, n: usize) -> T { if n % 2 == 0 { !xor } else { xor } }
706
707macro_rules! closed {
709 (@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)); };
710 (@one $r:ident, $t:ty, $unit:expr, $step:expr, $scan:ident, $slice:expr) => {
711 impl Reduction<$t> for $r {
712 type Total = $t;
713 const COMMUTATIVE: bool = true;
714 #[inline(always)]
715 fn unit() -> $t { $unit }
716 #[inline(always)]
717 fn step(acc: $t, x: $t) -> $t { ($step)(acc, x) }
718 #[inline(always)]
719 fn merge(a: $t, b: $t) -> ($t, bool) { (($step)(a, b), true) }
720 #[inline]
721 fn slice(x: &[$t]) -> ($t, bool) { (($slice)(x), true) }
722 #[inline]
723 fn scan(x: &[$t], start: Option<$t>, out: &mut Vec<$t>) -> bool {
724 closed!(@$scan x, start, out, $t, $step);
725 true
726 }
727 #[inline]
728 fn rows(data: &[$t], columns: usize, out: &mut [$t]) -> bool { rows_side_by_side::<$t, Self>(data, columns, out) }
729 }
730 };
731 (@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) };
732 (@grouped $x:ident, $start:ident, $out:ident, $t:ty, $step:expr) => { grouped($x, $start, $out, |v| v, $step, |t| t, true) };
733 (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())); )+};
734 ($($t:ty),+) => {$(
735 closed!(@one Max, $t, <$t>::MIN, Ord::max, records);
736 closed!(@one Min, $t, <$t>::MAX, Ord::min, records);
737 closed!(@one And, $t, !0, |a: $t, b: $t| a & b, grouped);
738 closed!(@one Or, $t, 0, |a: $t, b: $t| a | b, grouped);
739 closed!(@one Xor, $t, 0, |a: $t, b: $t| a ^ b, grouped);
740 )+};
741}
742closed!(i16, i32, i64, u8, u16, u32, u64);
743closed!(signed i16, i32, i64);
744impl Reduction<u8> for Xnor {
746 type Total = i16;
747 const COMMUTATIVE: bool = true;
748 #[inline(always)]
749 fn unit() -> i16 { !0 }
750 #[inline(always)]
751 fn step(acc: i16, x: u8) -> i16 { !(acc ^ i16::from(x)) }
752 #[inline(always)]
753 fn merge(a: i16, b: i16) -> (i16, bool) { (!(a ^ b), true) }
754 #[inline]
755 fn slice(x: &[u8]) -> (i16, bool) { (xnor_of(x.iter().fold(0i16, |a, &b| a ^ i16::from(b)), x.len()), true) }
756 #[inline]
757 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 }
758 #[inline]
759 fn rows(data: &[u8], columns: usize, out: &mut [i16]) -> bool { rows_side_by_side::<u8, Self>(data, columns, out) }
760}
761macro_rules! booleans {
763 ($($($r:ident)|+: $unit:expr, |$a:ident, $b:ident| $step:expr, $table:expr;)+) => {$($(
764 impl Reduction<bool> for $r {
765 type Total = bool;
766 const COMMUTATIVE: bool = true;
767 #[inline(always)]
768 fn unit() -> bool { $unit }
769 #[inline(always)]
770 fn step($a: bool, $b: bool) -> bool { $step }
771 #[inline(always)]
772 fn merge(a: bool, b: bool) -> (bool, bool) { (<Self as Reduction<bool>>::step(a, b), true) }
773 #[inline]
774 fn slice(x: &[bool]) -> (bool, bool) { (folded(x, $unit, <Self as Reduction<bool>>::step, <Self as Reduction<bool>>::step), true) }
775 #[inline]
776 fn scan(x: &[bool], start: Option<bool>, out: &mut Vec<bool>) -> bool {
777 Logical($table).scan(x, start, out);
778 true
779 }
780 #[inline]
781 fn short_lanes(data: &[bool], len: usize, start: Option<bool>, out: &mut Vec<bool>) -> Option<bool> {
782 (1..=64).contains(&len).then(|| { Logical($table).short_lanes(data, len, start, out); true })
783 }
784 }
785 )+)+};
786}
787booleans! {
788 Max | Or: false, |a, b| a | b, [[false, true], [true, true]];
789 Min | And: true, |a, b| a & b, [[false, false], [false, true]];
790 Xor: false, |a, b| a ^ b, [[false, true], [true, false]];
791 Xnor: true, |a, b| a == b, [[true, false], [false, true]];
792}
793
794pub(crate) fn records<T: Copy + PartialEq>(
796 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) {
797 let Some(&first) = x.first() else { return };
798 let first = from(first);
799 let (mut total, mut at, mut changes) = (start.map_or(first, |s| step(s, first)), 0, 0);
800 while at < x.len() {
801 let test = changes < 4;
802 let block = &x[at..x.len().min(at + if test { 256 } else { 256 << (changes - 4).min(6) })];
803 at += block.len();
804 if test && keeps(block, total) {
805 out.resize(out.len() + block.len(), total);
806 changes = 0;
807 } else {
808 grouped(block, Some(total), out, from, step, |t| t, false);
809 let last = out[out.len() - 1];
810 (total, changes) = (last, if last == total { 0 } else { changes + 1 });
811 }
812 }
813}
814
815#[inline(always)]
817pub(crate) fn below<F: PartialOrd + Copy>(block: &[F], total: F) -> bool { block.iter().fold(true, |k, &v| k & (v < total)) }
818
819#[inline(always)]
821pub(crate) fn above<F: PartialOrd + Copy>(block: &[F], total: F) -> bool { block.iter().fold(true, |k, &v| k & (v > total)) }
822
823const PIECE: usize = 1 << 18;
825const ROWS: usize = 16;
827
828fn piece_size(inner: usize) -> usize { PIECE.max(ROWS * inner).next_multiple_of(CHUNK) }
830
831#[inline(always)]
833fn one_of<T, R: Reduction<T>>(x: T, ok: &mut bool) -> R::Total {
834 let (t, fits) = R::one(x);
835 *ok &= fits;
836 t
837}
838
839#[inline(always)]
841fn steps_of<T: Copy, R: Reduction<T>>(totals: &mut [R::Total], x: &[T]) -> bool {
842 let (n, mut fits) = (totals.len().min(x.len()), true);
843 for k in 0..n { let (t, f) = R::checked(totals[k], x[k]); (totals[k], fits) = (t, fits & f) }
844 fits
845}
846
847#[inline(always)]
849fn started<T: Copy, R: Reduction<T>>(out: &mut Vec<R::Total>, rows: &[&[T]]) -> bool {
850 #[inline(always)]
851 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) }
852 let mut fits = true;
853 match *rows {
854 [] => {}
855 [a] => out.extend(a.iter().map(|&x| {
856 let (t, f) = R::one(x);
857 fits &= f;
858 t
859 })),
860 [a, b] => out.extend(a.iter().zip(b).map(|(&x, &y)| {
861 let (t, f) = step::<T, R>(R::one(x), y);
862 fits &= f;
863 t
864 })),
865 [a, b, c, ..] => out.extend(a.iter().zip(b).zip(c).map(|((&x, &y), &z)| {
866 let (t, f) = step::<T, R>(step::<T, R>(R::one(x), y), z);
867 fits &= f;
868 t
869 })),
870 }
871 fits
872}
873
874#[inline(always)]
876fn rows_folded<'a, T: Copy + 'a, R: Reduction<T>>(out: &mut Vec<R::Total>, width: usize, row: impl Fn(usize) -> &'a [T]) -> bool {
877 let at = out.len();
878 let mut fits = match width {
879 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)])
880 };
881 for i in 3..width { fits &= steps_of::<T, R>(&mut out[at..], row(i)); }
882 fits
883}
884
885#[inline(always)]
887fn merges_of<T, R: Reduction<T>>(totals: &mut [R::Total], parts: &[R::Total]) -> bool {
888 let (n, mut fits) = (totals.len().min(parts.len()), true);
889 for k in 0..n { let (t, f) = R::merge(totals[k], parts[k]); (totals[k], fits) = (t, fits & f) }
890 fits
891}
892
893#[inline(always)]
895fn merge_of<T, R: Reduction<T>>(a: R::Total, b: R::Total, ok: &mut bool) -> R::Total {
896 let (t, fits) = R::merge(a, b);
897 *ok &= fits;
898 t
899}
900
901#[inline(always)]
903fn chain_of<T, R: Reduction<T>>(a: R::Total, b: R::Total, ok: &mut bool) -> R::Total {
904 let (t, fits) = R::chain(a, b);
905 *ok &= fits;
906 t
907}
908
909const SHORT_LANE: usize = 64;
911
912const LANES: usize = 8;
914
915#[inline(always)]
917fn side_by_side<const B: usize, T: Copy, S: Copy>(
918 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)
919) -> bool {
920 let mut fits = true;
921 for (g, group) in data.chunks_exact(B * len).enumerate() {
922 let at = g * B * len;
923 let mut ok = [true; B];
924 let mut t: [S; B] = std::array::from_fn(|b| {
925 let (t, f) = first(group[b * len]);
926 ok[b] = f;
927 each(at + b * len, t);
928 t
929 });
930 for i in 1..len {
931 for b in 0..B {
932 let (s, f) = step(t[b], unsafe { *group.get_unchecked(b * len + i) });
934 (t[b], ok[b]) = (s, ok[b] & f);
935 each(at + b * len + i, s);
936 }
937 }
938 for (b, &t) in t.iter().enumerate() { last(g * B + b, t) }
939 fits &= ok.into_iter().all(|f| f);
940 }
941 fits
942}
943
944pub trait Term: Copy {
946 fn term<C: Tally>(self) -> C;
948 #[inline]
950 fn terms<C: Tally>(items: &[Self]) -> C { items.iter().fold(C::default(), |s, &x| s.add(x.term())) }
951}
952impl Term for bool {
953 #[inline(always)]
954 fn term<C: Tally>(self) -> C { C::low(self as i64) }
955 #[inline]
957 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()))) }
958}
959macro_rules! terms {
960 ($($t:ty),+) => {$(
961 impl Term for $t {
962 #[inline(always)]
963 fn term<C: Tally>(self) -> C { C::low(self as i64) }
964 }
965 )+};
966}
967terms!(i16, i32, i64, u8);
968
969pub trait Tally: Copy + Default + Send + Sync {
971 fn add(self, y: Self) -> Self;
973 fn low(n: i64) -> Self;
975}
976macro_rules! tallies {
977 ($($t:ty),+) => {$(
978 impl Tally for $t {
979 #[inline(always)]
980 fn add(self, y: $t) -> $t { self.wrapping_add(y) }
981 #[inline(always)]
982 fn low(n: i64) -> $t { n as $t }
983 }
984 )+};
985}
986tallies!(u8, i16, i32, i64);
987
988#[derive(Clone, Copy, Debug)]
990pub struct SumAs<C>(std::marker::PhantomData<C>);
991impl<T: Term, C: Tally> Reduction<T> for SumAs<C> {
992 type Total = C;
993 const COMMUTATIVE: bool = true;
994 #[inline(always)]
995 fn unit() -> C { C::default() }
996 #[inline(always)]
997 fn step(acc: C, x: T) -> C { acc.add(x.term()) }
998 #[inline(always)]
999 fn merge(a: C, b: C) -> (C, bool) { (a.add(b), true) }
1000 #[inline]
1001 fn slice(x: &[T]) -> (C, bool) { (T::terms(x), true) }
1002 #[inline]
1003 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 }
1004}
1005
1006#[derive(Clone, Copy, Debug)]
1008pub struct Squares;
1009
1010#[inline(always)]
1012fn add_square<T: Scalar>(s: f64, x: T) -> f64 {
1013 let re = x.re().wide();
1014 let s = re.mul_acc(re, s);
1015 if T::COMPLEX { let im = x.im().wide(); im.mul_acc(im, s) } else { s }
1016}
1017
1018impl<T: Scalar> Reduction<T> for Squares {
1019 type Total = f64;
1020 const COMMUTATIVE: bool = true;
1021 #[inline(always)]
1022 fn unit() -> f64 { 0.0 }
1023 #[inline(always)]
1024 fn step(acc: f64, x: T) -> f64 { add_square(acc, x) }
1025 #[inline(always)]
1026 fn merge(a: f64, b: f64) -> (f64, bool) { (a + b, true) }
1027 #[inline]
1028 fn slice(x: &[T]) -> (f64, bool) { (folded(x, 0.0, add_square, |a, b| a + b), true) }
1029 #[inline]
1030 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 }
1031}
1032
1033#[inline(always)]
1035fn 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 {
1036 let (pairs, rest) = x.as_chunks::<2>();
1037 let total = folded(pairs, unit, |s, [a, b]| pair(s, a, b), merge);
1038 rest.first().map_or(total, |&z| last(total, z))
1039}
1040
1041#[inline(always)]
1043fn quotient<T: Copy, R: Native>(
1044 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) {
1045 let step = move |(p, low): (T, R), q: T| { let r = times(p, q); (r, low.min(size(r))) };
1046 let (total, low) = pairwise(x, (one, R::one()), |s, a, b| step(s, divide(a, b)), move |s, (q, l)| {
1047 let (r, m) = step(s, q);
1048 (r, m.min(l))
1049 }, step);
1050 (total, low >= R::min_positive_value() && finite(total))
1051}
1052
1053#[inline(always)]
1055fn 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 {
1056 let merge = &merge;
1057 running_results(x, unit, add, merge, move |parts, rest| {
1058 let ((pairs, _), (tail, last)) = (parts.as_chunks::<2>(), rest.as_chunks::<2>());
1059 let (even, odd) = pairs.iter().fold((unit, unit), |(e, o), &[p, q]| (merge(e, p), merge(o, q)));
1060 let (even, odd) = tail.iter().fold((even, odd), |(e, o), &[p, q]| (add(e, p), add(o, q)));
1061 last.first().map_or(even, |&z| add(even, z)) - odd
1062 })
1063}
1064
1065fn checked_steps<T: Copy, R: Reduction<T>>(x: &[T], start: Option<R::Total>, out: &mut Vec<R::Total>) -> bool {
1067 let (mut total, mut fits) = (start.unwrap_or_else(R::identity), true);
1068 out.extend(x.iter().map(|&v| {
1069 let ok;
1070 (total, ok) = R::checked(total, v);
1071 fits &= ok;
1072 total
1073 }));
1074 fits
1075}
1076
1077#[derive(Clone, Copy, Debug, Default)]
1079pub struct Alternated<R>(std::marker::PhantomData<R>);
1080
1081pub trait Inverse<T>: Reduction<T> {
1083 fn inverse(acc: Self::Total, x: T) -> (Self::Total, bool);
1085 fn undo(a: Self::Total, b: Self::Total) -> (Self::Total, bool);
1087 #[inline(always)]
1089 fn fits(_: Self::Total) -> bool { true }
1090 fn alternating(x: &[T]) -> (Self::Total, bool);
1092}
1093
1094impl<T: Copy, R: Inverse<T>> Reduction<T> for Alternated<R> {
1095 type Total = (R::Total, bool);
1096 #[inline(always)]
1097 fn unit() -> Self::Total { (R::unit(), false) }
1098 #[inline(always)]
1099 fn identity() -> Self::Total { (R::identity(), false) }
1100 #[inline(always)]
1101 fn step((s, odd): Self::Total, x: T) -> Self::Total { (if odd { R::inverse(s, x).0 } else { R::step(s, x) }, !odd) }
1102 #[inline(always)]
1103 fn checked((s, odd): Self::Total, x: T) -> (Self::Total, bool) {
1104 let (t, ok) = if odd { R::inverse(s, x) } else { R::checked(s, x) };
1105 ((t, !odd), ok && R::fits(t))
1106 }
1107 #[inline(always)]
1108 fn merge((a, odd): Self::Total, (b, more): Self::Total) -> (Self::Total, bool) {
1109 let (t, ok) = if odd { R::undo(a, b) } else { R::merge(a, b) };
1110 ((t, odd ^ more), ok && R::fits(t))
1111 }
1112 #[inline]
1113 fn slice(x: &[T]) -> (Self::Total, bool) { let (t, ok) = R::alternating(x); ((t, x.len() % 2 == 1), ok) }
1114}
1115
1116impl<F: Native> Inverse<F> for Sum {
1117 #[inline(always)]
1118 fn inverse(acc: F, x: F) -> (F, bool) { (acc.algebraic_add(-x), true) }
1119 #[inline(always)]
1120 fn undo(a: F, b: F) -> (F, bool) { (a.algebraic_add(-b), true) }
1121 #[inline]
1122 fn alternating(x: &[F]) -> (F, bool) { (alternated(x, F::zero(), |s, v| s + v, |a, b| a + b), true) }
1123}
1124
1125impl<F: Native> Inverse<Complex<F>> for Sum {
1126 #[inline(always)]
1127 fn inverse(acc: Complex<F>, x: Complex<F>) -> (Complex<F>, bool) { (<Sum as Reduction<Complex<F>>>::step(acc, -x), true) }
1128 #[inline(always)]
1129 fn undo(a: Complex<F>, b: Complex<F>) -> (Complex<F>, bool) { (<Sum as Reduction<Complex<F>>>::step(a, -b), true) }
1130 #[inline]
1131 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) }
1132}
1133
1134macro_rules! alternated_integers {
1136 ($($t:ty => $acc:ty, $run:expr);+) => {$(
1137 impl Inverse<$t> for Sum {
1138 #[inline(always)]
1139 fn inverse(acc: i64, x: $t) -> (i64, bool) { let (t, over) = acc.overflowing_sub(i64::from(x)); (t, !over) }
1140 #[inline(always)]
1141 fn undo(a: i64, b: i64) -> (i64, bool) { let (t, over) = a.overflowing_sub(b); (t, !over) }
1142 #[inline]
1143 fn alternating(x: &[$t]) -> (i64, bool) {
1144 let run = |c: &[$t]| i128::from(alternated(c, 0 as $acc, |s, v| s + <$acc>::from(v), |s, t| s + t));
1145 let total: i128 = x.chunks($run).map(run).sum();
1146 (total as i64, i64::try_from(total).is_ok())
1147 }
1148 }
1149 )+};
1150}
1151alternated_integers!(i16 => i32, 1 << 15; u8 => i32, 1 << 15; i32 => i64, 1 << 31; i64 => i128, 1 << 31);
1152
1153#[inline(always)]
1155fn normal<F: Native>(q: F) -> bool { q.abs() >= F::min_positive_value() && q.is_finite() }
1156
1157#[derive(Clone, Copy, Debug)]
1159pub struct Scaled<F> { m: F, e: F, normal: bool }
1160
1161impl<F: Native> Scaled<F> {
1162 #[inline(always)]
1164 fn of(p: F, e: F, normal: bool) -> Self {
1165 let (lo, hi) = if F::WIDE { (1.4916681462400413e-154, 6.703903964971299e153) } else { (1.0842021724855044e-19, 9.223372036854776e18) };
1166 if p.abs() >= F::narrow(lo) && p.abs() <= F::narrow(hi) { Self { m: p, e, normal } } else { Self::rescaled(p, e, normal) }
1167 }
1168 #[inline]
1170 fn rescaled(p: F, e: F, normal: bool) -> Self {
1171 let (m, f) = p.abs().split();
1172 Self { m: if p < F::zero() { -m } else { m }, e: e + f, normal: normal & p.abs().positive_normal() }
1173 }
1174 #[inline(always)]
1176 fn times(self, q: F) -> Self { Self::of(self.m * q, self.e, self.normal) }
1177 #[inline(always)]
1179 fn over(self, q: F) -> Self { Self::of(self.m / q, self.e, self.normal) }
1180 #[inline(always)]
1182 fn joined(self, b: Self) -> Self { Self::of(self.m * b.m, self.e + b.e, self.normal & b.normal) }
1183 #[inline(always)]
1185 fn undone(self, b: Self) -> Self { Self::of(self.m / b.m, self.e - b.e, self.normal & b.normal) }
1186 #[inline]
1188 pub fn value(self) -> F {
1189 let Self { m, e, .. } = Self::rescaled(self.m, self.e, true);
1190 let far = F::narrow(if F::WIDE { 2040.0 } else { 250.0 });
1191 crate::float::pow2_times::<F, { crate::FUSED }>(m, e.max(-far).min(far))
1192 }
1193}
1194
1195#[derive(Clone, Copy, Debug, Default)]
1197pub struct ScaledProduct;
1198
1199impl<F: Native> Reduction<F> for ScaledProduct {
1200 type Total = Scaled<F>;
1201 const COMMUTATIVE: bool = true;
1202 #[inline(always)]
1203 fn unit() -> Scaled<F> { Scaled { m: F::one(), e: F::zero(), normal: true } }
1204 #[inline(always)]
1205 fn step(acc: Scaled<F>, x: F) -> Scaled<F> { acc.times(x) }
1206 #[inline(always)]
1207 fn checked(acc: Scaled<F>, x: F) -> (Scaled<F>, bool) { let t = acc.times(x); (t, t.normal) }
1208 #[inline(always)]
1209 fn merge(a: Scaled<F>, b: Scaled<F>) -> (Scaled<F>, bool) { let t = a.joined(b); (t, t.normal) }
1210}
1211
1212impl<F: Native> Inverse<F> for ScaledProduct {
1213 #[inline(always)]
1214 fn inverse(acc: Scaled<F>, x: F) -> (Scaled<F>, bool) { let t = acc.over(x); (t, t.normal) }
1215 #[inline(always)]
1216 fn undo(a: Scaled<F>, b: Scaled<F>) -> (Scaled<F>, bool) { let t = a.undone(b); (t, t.normal) }
1217 #[inline]
1218 fn alternating(x: &[F]) -> (Scaled<F>, bool) {
1219 let (q, fits) = quotient(x, F::one(), |a, b| a / b, F::algebraic_mul, F::abs, F::is_finite);
1221 let t = if fits { Scaled::of(q, F::zero(), true) }
1222 else { pairwise(x, <Self as Reduction<F>>::unit(), |s, a, b| s.times(a / b), Scaled::joined, Scaled::times) };
1223 (t, t.normal)
1224 }
1225}
1226
1227#[inline(always)]
1229fn complex_times<F: Native>(p: Complex<F>, q: Complex<F>) -> Complex<F> {
1230 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)))
1231}
1232
1233#[inline(always)]
1235fn complex_size<F: Native>(p: Complex<F>) -> F { p.re.abs().max(p.im.abs()) }
1236
1237impl<F: Native> Inverse<Complex<F>> for Product {
1239 #[inline(always)]
1240 fn inverse(acc: Complex<F>, x: Complex<F>) -> (Complex<F>, bool) { (crate::divide(acc, x), true) }
1241 #[inline(always)]
1242 fn undo(a: Complex<F>, b: Complex<F>) -> (Complex<F>, bool) { (crate::divide(a, b), true) }
1243 #[inline(always)]
1244 fn fits(q: Complex<F>) -> bool { normal(complex_size(q)) }
1245 #[inline]
1246 fn alternating(x: &[Complex<F>]) -> (Complex<F>, bool) {
1247 quotient(x, Complex::new(F::one(), F::zero()), crate::divide, complex_times, complex_size, |p| p.re.is_finite() & p.im.is_finite())
1248 }
1249}
1250
1251#[macro_export]
1253macro_rules! with_reduce {
1254 ($r:expr, |$t:ident| $body:expr) => {
1255 match $r {
1256 $crate::reduce::Reduce::Sum => { type $t = $crate::reduce::Sum; $body }
1257 $crate::reduce::Reduce::Product => { type $t = $crate::reduce::Product; $body }
1258 $crate::reduce::Reduce::Max => { type $t = $crate::reduce::Max; $body }
1259 $crate::reduce::Reduce::Min => { type $t = $crate::reduce::Min; $body }
1260 $crate::reduce::Reduce::And => { type $t = $crate::reduce::And; $body }
1261 $crate::reduce::Reduce::Or => { type $t = $crate::reduce::Or; $body }
1262 $crate::reduce::Reduce::Xor => { type $t = $crate::reduce::Xor; $body }
1263 $crate::reduce::Reduce::Xnor => { type $t = $crate::reduce::Xnor; $body }
1264 }
1265 };
1266}
1267
1268#[doc(hidden)]
1270#[macro_export]
1271macro_rules! native_reduction {
1272 (@inverse $kind:expr, $dtype:expr, |$R:ident, $T:ident| $body:expr) => {{
1273 use $crate::buffer::{DType, FloatWidth};
1274 use $crate::reduce::{Alternated, Product, Reduce, ScaledProduct, Sum};
1275 match ($dtype, $kind) {
1276 (DType::Integer(w), Reduce::Sum) => with_width!(w, I => $crate::native_reduction!(@at Alternated<Sum>, I, |$R, $T| $body)),
1277 (DType::Float(FloatWidth::F16 | FloatWidth::F32), Reduce::Sum) => $crate::native_reduction!(@at Alternated<Sum>, f32, |$R, $T| $body),
1278 (DType::Float(FloatWidth::F64), Reduce::Sum) => $crate::native_reduction!(@at Alternated<Sum>, f64, |$R, $T| $body),
1279 (DType::Float(FloatWidth::F16 | FloatWidth::F32), Reduce::Product) => $crate::native_reduction!(@at Alternated<ScaledProduct>, f32, |$R, $T| $body),
1280 (DType::Float(FloatWidth::F64), Reduce::Product) => $crate::native_reduction!(@at Alternated<ScaledProduct>, f64, |$R, $T| $body),
1281 (DType::Complex(w), Reduce::Sum) => with_complex_width!(w, C => $crate::native_reduction!(@at Alternated<Sum>, C, |$R, $T| $body)),
1282 (DType::Complex(w), Reduce::Product) => with_complex_width!(w, C => $crate::native_reduction!(@at Alternated<Product>, C, |$R, $T| $body)),
1283 _ => None,
1284 }
1285 }};
1286 (@at $r:ty, $t:ty, |$R:ident, $T:ident| $body:expr) => {{ type $R = $r; type $T = $t; $body }};
1287 ($kind:expr, $dtype:expr, $checked:expr, |$R:ident, $T:ident| $body:expr) => {{
1288 use $crate::buffer::DType;
1289 use $crate::reduce::Reduce;
1290 match ($dtype, $kind) {
1291 (DType::Boolean, Reduce::Sum) => $crate::native_reduction!(@at $crate::reduce::SumAs<i64>, bool, |$R, $T| $body),
1292 (DType::Boolean, Reduce::Product | Reduce::And | Reduce::Min) => $crate::native_reduction!(@at $crate::reduce::And, bool, |$R, $T| $body),
1293 (DType::Boolean, Reduce::Or | Reduce::Max) => $crate::native_reduction!(@at $crate::reduce::Or, bool, |$R, $T| $body),
1294 (DType::Boolean, Reduce::Xor) => $crate::native_reduction!(@at $crate::reduce::Xor, bool, |$R, $T| $body),
1295 (DType::Boolean, Reduce::Xnor) => $crate::native_reduction!(@at $crate::reduce::Xnor, bool, |$R, $T| $body),
1296 (DType::Integer(w), Reduce::Sum) if !$checked => with_width!(w, I => $crate::native_reduction!(@at $crate::Wrapped<$crate::reduce::Sum>, I, |$R, $T| $body)),
1297 (DType::Integer(w), Reduce::Sum) => with_width!(w, I => $crate::native_reduction!(@at $crate::reduce::Sum, I, |$R, $T| $body)),
1298 (DType::Integer(w), Reduce::Product) if !$checked => with_width!(w, I => $crate::native_reduction!(@at $crate::Wrapped<$crate::reduce::Product>, I, |$R, $T| $body)),
1299 (DType::Integer(w), Reduce::Product) => with_width!(w, I => $crate::native_reduction!(@at $crate::reduce::Product, I, |$R, $T| $body)),
1300 (DType::Integer(w), Reduce::Max) => with_width!(w, I => $crate::native_reduction!(@at $crate::reduce::Max, I, |$R, $T| $body)),
1301 (DType::Integer(w), Reduce::Min) => with_width!(w, I => $crate::native_reduction!(@at $crate::reduce::Min, I, |$R, $T| $body)),
1302 (DType::Integer(w), Reduce::And) => with_width!(w, I => $crate::native_reduction!(@at $crate::reduce::And, I, |$R, $T| $body)),
1303 (DType::Integer(w), Reduce::Or) => with_width!(w, I => $crate::native_reduction!(@at $crate::reduce::Or, I, |$R, $T| $body)),
1304 (DType::Integer(w), Reduce::Xor) => with_width!(w, I => $crate::native_reduction!(@at $crate::reduce::Xor, I, |$R, $T| $body)),
1305 (DType::Integer(w), Reduce::Xnor) => with_width!(w, I => $crate::native_reduction!(@at $crate::reduce::Xnor, I, |$R, $T| $body)),
1306 (DType::Float(w), Reduce::Sum) => with_float_width!(w, F => $crate::native_reduction!(@at $crate::reduce::Sum, F, |$R, $T| $body)),
1307 (DType::Float(w), Reduce::Product) => with_float_width!(w, F => $crate::native_reduction!(@at $crate::reduce::Product, F, |$R, $T| $body)),
1308 (DType::Float(w), Reduce::Max) => with_float_width!(w, F => $crate::native_reduction!(@at $crate::reduce::Max, F, |$R, $T| $body)),
1309 (DType::Float(w), Reduce::Min) => with_float_width!(w, F => $crate::native_reduction!(@at $crate::reduce::Min, F, |$R, $T| $body)),
1310 (DType::Complex(w), Reduce::Sum) => with_complex_width!(w, C => $crate::native_reduction!(@at $crate::reduce::Sum, C, |$R, $T| $body)),
1311 (DType::Complex(w), Reduce::Product) => with_complex_width!(w, C => $crate::native_reduction!(@at $crate::reduce::Product, C, |$R, $T| $body)),
1312 _ => None,
1313 }
1314 }};
1315}
1316pub use crate::native_reduction;
1317
1318macro_rules! wrapped_products {
1319 ($($t:ty),+) => {$(
1320 impl Reduction<$t> for crate::Wrapped<Product> {
1321 type Total = i64;
1322 const COMMUTATIVE: bool = true;
1323 fn unit() -> i64 { 1 }
1324 fn step(total: i64, x: $t) -> i64 { total.wrapping_mul(x.into()) }
1325 fn merge(x: i64, y: i64) -> (i64, bool) { (x.wrapping_mul(y), true) }
1326 fn slice(x: &[$t]) -> (i64, bool) { (folded(x, 1i64, |total, x| total.wrapping_mul(x.into()), i64::wrapping_mul), true) }
1327 fn scan(x: &[$t], start: Option<i64>, out: &mut Vec<i64>) -> bool {
1328 grouped(x, start, out, |x| x.into(), i64::wrapping_mul, |x| x, true);
1329 true
1330 }
1331 }
1332 )+};
1333}
1334wrapped_products!(u8, i16, i32, i64);
1335
1336mod along;
1337mod pieces;
1338mod run;
1339mod scans;
1340mod sliding;
1341pub use {along::*, pieces::*, run::*, scans::*, sliding::*};