1use super::*;
2
3#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
5pub enum Reduce { Sum, Product, Max, Min, And, Or, Xor, Xnor }
6
7impl Reduce {
8 pub fn stage(self, item: crate::buffer::DType, exact: bool) -> bool {
10 use crate::buffer::{DType, FloatWidth};
11 match item {
12 DType::Float(FloatWidth::F32 | FloatWidth::F64) => matches!(self, Self::Sum | Self::Product | Self::Max | Self::Min),
13 DType::Integer(_) => matches!(self, Self::Max | Self::Min) || exact && self == Self::Sum,
14 DType::Boolean => matches!(self, Self::Sum | Self::Min | Self::Max | Self::Xor | Self::Xnor), _ => false
15 }
16 }
17
18 pub fn seeds(self, items: crate::buffer::DType, seed: crate::buffer::DType, result: crate::buffer::DType) -> bool {
20 use crate::buffer::DType;
21 match (items, seed, result) {
22 (DType::Float(a), DType::Float(_), DType::Float(b)) => a == b,
23 (DType::Integer(_), DType::Boolean | DType::Integer(_), DType::Integer(_)) => match self {
24 Self::Max | Self::Min => true, Self::Sum => items == result, _ => false
25 },
26 (DType::Boolean, DType::Boolean | DType::Integer(_), DType::Integer(_)) => self == Self::Sum,
27 _ => false,
28 }
29 }
30
31 pub fn compiles(self) -> bool { matches!(self, Self::Sum | Self::Product | Self::Max | Self::Min) }
33}
34impl Reduce {
35 pub fn total_dtype(self, variant: crate::domain::Variant) -> Option<crate::buffer::DType> {
36 native_reduction!(self, variant.dtype, variant.fails_on_overflow(), |R, T| Some(<<R as Reduction<T>>::Total as Element>::DTYPE))
37 }
38 pub fn run_pieces(self, variant: crate::domain::Variant, dims: [usize; 3], seed: Seed<'_>, pieces: &Pieces<'_>, output: &mut Buffer) -> Result<(), Error> {
39 use crate::buffer::{DType, Width};
41 if self == Reduce::Sum && variant.dtype == DType::Boolean && matches!(seed, Seed::Identity)
42 && let DType::Integer(w @ (Width::I16 | Width::I32)) = output.dtype() && crate::domain::Range::of_dtype(output.dtype()).hi >= dims[1] as f64 {
43 return if w == Width::I16 { counted::<i16>(dims, variant.checked, pieces, output) }
44 else { counted::<i32>(dims, variant.checked, pieces, output) };
45 }
46 native_reduction!(self, variant.dtype, variant.fails_on_overflow(), |R, T| Some(
47 native_pieces::<T, R>(dims, variant.checked, pieces).and_then(|result| finish_fold::<T, R>(result, seed, variant.checked, output))
48 )).unwrap_or_else(|| Err(Error::new(Kind::Domain, "reduction does not support this dtype")))
49 }
50 pub fn run_read(
51 self, variant: crate::domain::Variant, shape: &[usize], axis: Option<usize>, seed: Seed<'_>, input: Read<'_>, output: &mut Buffer
52 ) -> Result<Option<Vec<Steps>>, Error> {
53 native_reduction!(self, variant.dtype, variant.fails_on_overflow(), |R, T| Some(native_read::<T, R>(shape, axis, variant.checked, input).and_then(
54 |(buffer, fits, order)| {
55 let (buffer, order) = if seed.read().is_some() { (buffer.ordered(order), None) } else { (buffer, order) };
57 finish_fold::<T, R>((buffer, fits), seed, variant.checked, output).map(|()| order)
58 }
59 ))).unwrap_or_else(|| Err(Error::new(Kind::Domain, "reduction does not support this dtype")))
60 }
61 pub fn run_sliding_2d(
63 self, variant: crate::domain::Variant, dims: [usize; 5], widths: [usize; 2], every: [usize; 2], input: View<'_>, output: &mut Buffer
64 ) -> Result<(), Error> {
65 native_reduction!(self, variant.dtype, variant.fails_on_overflow(), |R, T| Some(items::<T>(input).and_then(|data| {
66 let folds = sliding_2d::<T, R>(data, dims, widths, every).ok_or_else(|| Error::new(Kind::Domain, "windows too wide to fold directly"))?;
67 totals::<T, R>(folds, variant.checked, output)
68 }))).unwrap_or_else(|| Err(Error::new(Kind::Domain, "reduction does not support this dtype")))
69 }
70 pub fn run_groups(self, variant: crate::domain::Variant, groups: &impl crate::search::Grouping, input: View<'_>, output: &mut Buffer) -> Result<(), Error> {
72 native_reduction!(self, variant.dtype, variant.fails_on_overflow(), |R, T| Some(items::<T>(input).and_then(|data| totals::<T, R>(
73 groups.fold::<T, R>(data), variant.checked, output
74 )))).unwrap_or_else(|| Err(Error::new(Kind::Domain, "reduction does not support this dtype")))
75 }
76 #[inline(never)]
77 pub fn inverse_pieces(self, variant: crate::domain::Variant, dims: [usize; 3], pieces: &Pieces<'_>) -> Result<Option<Buffer>, Error> {
78 inverse_result(
79 native_reduction!(@inverse self, variant.dtype, |R, T| Some(native_pieces::<T, R>(dims, variant.checked, pieces))), dims.into_iter().product())
80 }
81 #[inline(never)]
82 pub fn inverse_read(self, variant: crate::domain::Variant, shape: &[usize], axis: usize, input: Read<'_>) -> Result<Option<Buffer>, Error> {
83 inverse_result(
84 native_reduction!(@inverse self, variant.dtype, |R, T| Some(native_read::<T, R>(shape, Some(axis), variant.checked, input).map(|(buffer, fits, order)| (buffer.ordered(order), fits)))),
85 crate::layout::count(shape))
86 }
87}
88
89fn inverse_result(result: Option<Result<(Buffer, bool), Error>>, items: usize) -> Result<Option<Buffer>, Error> {
90 let result = result.transpose()?.and_then(|(buffer, fits)| fits.then_some(buffer));
91 if result.is_some() { crate::events::taken(crate::events::Route::AlternatedFold, items); }
92 Ok(result)
93}
94
95fn native_read<T, R>(shape: &[usize], axis: Option<usize>, checked: bool, input: Read<'_>) -> Result<(Buffer, bool, Option<Vec<Steps>>), Error> where T: Cast,
96 R: Reduction<T>, R::Total: Totals {
97 let dims = match axis { Some(axis) => Lanes::new(shape, axis)?.dims(), None => [1, crate::layout::count(shape), 1] };
98 if let Some((view, mapping)) = input.positional(shape) && let Some(data) = T::of(view) {
99 if axis.is_none() && R::COMMUTATIVE && let Some(block) = mapping.block(shape) {
100 return crate::error::checked(|| {
101 let (dims, order) = block.fold(None);
102 let (totals, fits) = along::<T, R>(&data[block.range], dims);
103 Ok((R::Total::buffer_of(totals), fits, order))
104 });
105 }
106 if let Some(axis) = axis {
107 return crate::error::checked(|| {
108 let slide = if R::COMMUTATIVE {
109 let source = crate::structure::walked_source(&mapping.tables(shape), shape, data.len());
110 match fold_route(&mapping, shape, source.as_deref(), false, &[axis]) {
111 FoldRoute::Slide(s) => {
112 let data = &data[s.region.clone()];
113 let folded = match s.dims[..] {
114 [outer, len, inner] if s.every[0] == 1 => Some(sliding::<T, R>(data, [outer, len, inner], s.widths[0])),
115 [outer, len, inner] => {
116 crate::events::taken(crate::events::Route::SlidingFold, data.len());
117 Some(along::<T, R>(data, [outer, len, inner]))
118 }
119 [a, b, c, d, e] => sliding_2d::<T, R>(data, [a, b, c, d, e], [s.widths[0], s.widths[1]], [s.every[0], s.every[1]]),
120 _ => unreachable!("a slide folds one axis or two"),
121 };
122 folded.map(|(totals, fits)| (totals, fits, Some(s.tables)))
123 }
124 _ => None,
125 }
126 } else { None };
127 let (totals, fits, order) = slide.unwrap_or_else(|| along_mapped::<T, R>(data, &mapping, shape, axis));
128 let order = order.filter(|order| !Mapping::walk_in_order(0, order, totals.len()));
129 Ok((R::Total::buffer_of(totals), fits, order))
130 });
131 }
132 }
133 let pieces = |range: Range<usize>, emit: &mut dyn FnMut(usize, Piece<'_>)| {
134 let mut scratch = Buffer::with_capacity(T::DTYPE, 0);
135 let start = range.start;
136 let view = input.part_view(range, T::DTYPE, checked, &mut scratch)?;
137 emit(start, Piece::Items(view));
138 Ok(())
139 };
140 native_pieces::<T, R>(dims, checked, &pieces).map(|(buffer, fits)| (buffer, fits, None))
141}
142
143fn counted<C: Tally + Cast>(dims: [usize; 3], checked: bool, pieces: &Pieces<'_>, output: &mut Buffer) -> Result<(), Error> {
145 native_pieces::<bool, SumAs<C>>(dims, checked, pieces).and_then(|result| finish_fold::<bool, SumAs<C>>(result, Seed::Identity, checked, output))
146}
147
148fn native_pieces<T, R>(dims: [usize; 3], checked: bool, pieces: &Pieces<'_>) -> Result<(Buffer, bool), Error> where T: Cast, R: Reduction<T>, R::Total: Totals {
149 let [outer, len, inner] = dims;
150 let mut totals = vec![if len == 0 { R::unit() } else { R::identity() }; outer * inner];
151 crate::error::checked(|| {
152 let fits = fold_pieces::<T, R>(&mut totals, dims, &|range, emit| {
153 let result = pieces(range, &mut |start, piece| match piece {
154 Piece::Items(view) => match T::of(view) {
155 Some(items) => emit(start, Part::Items(items)),
156 None => match Buffer::cast_results(view, T::DTYPE, checked) {
157 Ok(values) => emit(start, Part::Items(T::of(values.view()).expect("converted fold piece has its requested dtype"))),
158 Err(error) => crate::error::record(error),
159 },
160 },
161 Piece::Total(total) => match R::Total::from_partial(total) {
162 Ok(value) => emit(start, Part::Total(total.items, value)), Err(error) => crate::error::record(error)
163 },
164 });
165 if let Err(error) = result { crate::error::record(error); }
166 });
167 Ok((R::Total::buffer_of(totals), fits))
168 })
169}
170
171fn finish_fold<T, R: Reduction<T>>((mut buffer, fits): (Buffer, bool), seed: Seed<'_>, checked: bool, output: &mut Buffer) -> Result<(), Error> where
172 R::Total: Cast {
173 crate::error::checked(|| {
174 if checked && !fits { return Err(Error::new(Kind::Overflow, "integer fold overflowed")); }
175 let totals = R::Total::of_mut(&mut buffer).expect("an associative total keeps its storage type");
176 if let Some(read) = seed.read() {
177 let values = seed_row::<R::Total>(&read, 0, totals.len(), checked, &mut Vec::new())?;
178 for (total, value) in totals.iter_mut().zip(values) {
179 let (result, fits) = if matches!(seed, Seed::Initial(_)) { R::merge(value, *total) } else { R::merge(*total, value) };
180 *total = result;
181 if checked && !fits { crate::error::record(Error::new(Kind::Overflow, "integer fold overflowed")); }
182 }
183 }
184 *output = buffer;
185 Ok(())
186 })
187}
188
189fn items<'a, T: Element>(view: View<'a>) -> Result<&'a [T], Error> { T::of(view).ok_or_else(|| Error::new(Kind::Domain, "fold input has the wrong dtype")) }
191
192fn totals<T, R: Reduction<T>>((totals, fits): (Vec<R::Total>, bool), checked: bool, output: &mut Buffer) -> Result<(), Error> where R::Total: Cast + Totals {
194 finish_fold::<T, R>((R::Total::buffer_of(totals), fits), Seed::Identity, checked, output)
195}
196
197impl Reduce {
198 pub fn run_scan(
199 self, variant: crate::domain::Variant, dims: [usize; 3], seed: Seed<'_>, direction: Direction, input: crate::agreement::Read<'_>,
200 output: &mut crate::buffer::Buffer) -> Result<(), Error> {
201 if self == Reduce::Sum && let Some(done) = bounded_scan(variant, dims, seed, direction, false, input, output) { return done; }
202 native_reduction!(self, variant.dtype, variant.fails_on_overflow(), |R, T| Some(native_scan::<T, R>(
203 dims, seed, direction, variant.checked, input, output
204 ))).unwrap_or_else(|| Err(Error::new(Kind::Domain, "scan does not support this dtype")))
205 }
206}
207
208fn native_scan<T, R>(
209 dims: [usize; 3], seed: Seed<'_>, direction: Direction, checked: bool, input: crate::agreement::Read<'_>, output: &mut crate::buffer::Buffer
210) -> Result<(), Error> where T: crate::buffer::Cast, R: Reduction<T>, R::Total: crate::buffer::Cast {
211 scan_read::<T, R::Total>(dims, seed, direction, checked, input, output, &|items, dims, seed, out| {
212 let fits = scan_along::<T, R>(items, dims, seed, out);
213 if checked && !fits { Err(Error::new(Kind::Overflow, "integer scan overflowed")) } else { Ok(()) }
214 })
215}
216
217pub fn negated_scan(
219 variant: crate::domain::Variant, dims: [usize; 3], seed: Seed<'_>, input: crate::agreement::Read<'_>, output: &mut crate::buffer::Buffer
220) -> Result<bool, Error> {
221 use crate::buffer::{DType, FloatWidth};
222 macro_rules! negated {
223 ($t:ty) => {
224 scan_read::<$t, $t>(dims, seed, Direction::Left, variant.checked, input, output, &|items, dims, seed, out| {
225 super::scans::negated_cells::<$t, Sum>(items, dims, seed, out, |x| -x);
226 Ok(())
227 })
228 };
229 }
230 match variant.dtype {
231 DType::Float(FloatWidth::F16 | FloatWidth::F32) => negated!(f32)?,
232 DType::Float(FloatWidth::F64) => negated!(f64)?,
233 DType::Complex(w) => with_complex_width!(w, C => negated!(C))?,
234 DType::Integer(_) | DType::Boolean => {
235 return bounded_scan(variant, dims, seed, Direction::Left, true, input, output).map_or(Ok(false), |done| done.map(|()| true));
236 }
237 _ => return Ok(false),
238 }
239 Ok(true)
240}
241
242pub fn bounded_scan(
244 variant: crate::domain::Variant, dims: [usize; 3], seed: Seed<'_>, direction: Direction, negate: bool, input: crate::agreement::Read<'_>,
245 output: &mut crate::buffer::Buffer) -> Option<Result<(), Error>> {
246 use crate::buffer::{DType, Width};
247 let DType::Integer(total @ (Width::I16 | Width::I32)) = output.dtype() else { return None };
248 let fits = |d: DType| d == DType::Boolean || matches!(d, DType::Integer(w) if total.join(w) == total);
249 let stored = input.plan.try_for_each_source(&mut |s, _| if input.sources.get(s).is_some_and(|b| fits(b.view.dtype())) { Ok(()) } else { Err(()) });
250 if !matches!(variant.dtype, DType::Boolean | DType::Integer(_)) || stored.is_err() { return None; }
251 let checked = variant.fails_on_overflow();
252 Some(if total == Width::I16 { totalled::<i16>(dims, seed, direction, checked, negate, input, output) }
253 else { totalled::<i32>(dims, seed, direction, checked, negate, input, output) })
254}
255
256fn totalled<C>(
258 dims: [usize; 3], seed: Seed<'_>, direction: Direction, checked: bool, negate: bool, input: crate::agreement::Read<'_>, output: &mut crate::buffer::Buffer
259) -> Result<(), Error> where C: crate::buffer::Cast + Tally + Term + Into<i64> + Ord + Default + std::ops::BitXor<Output = C> + std::ops::BitAnd<Output = C>
260 + std::ops::BitOr<Output = C> {
261 scan_read::<C, C>(dims, seed, direction, checked, input, output, &|items, [g, n, inner], seed, out| {
262 let from = out.len();
263 if negate { super::scans::negated_cells::<C, SumAs<C>>(items, [g, n, inner], seed, out, |x| C::low(-x.into())) }
264 else { scan_along::<C, SumAs<C>>(items, [g, n, inner], seed, out); }
265 if !checked { return Ok(()); }
266 let test = |p: C, x: C, c: C| if negate { (p ^ x) & (p ^ c) } else { (p ^ c) & (x ^ c) };
267 let along = |p: &[C], x: &[C], c: &[C]| p.iter().zip(x).zip(c).fold(C::default(), |b, ((&p, &x), &c)| b | test(p, x, c));
268 let mut bad = C::default();
269 for (x, c) in items.chunks_exact(n * inner).zip(out[from..].chunks_exact(n * inner)) {
270 bad = bad | along(&c[..(n - 1) * inner], &x[inner..], &c[inner..]);
271 if let Some(s) = seed { bad = bad | (0..inner).fold(C::default(), |b, i| b | test(s[if s.len() == 1 { 0 } else { i }], x[i], c[i])); }
272 }
273 if bad < C::default() { Err(Error::new(Kind::Overflow, "integer scan overflowed")) } else { Ok(()) }
274 })
275}
276
277pub fn seed_row<T: crate::buffer::Cast>(
279 read: &crate::agreement::Read<'_>, cell: usize, inner: usize, checked: bool, scratch: &mut Vec<T>) -> Result<Vec<T>, Error> {
280 let one = read.plan.len() == 1;
281 let arg = read.part::<T>(if one { 0..1 } else { cell * inner..(cell + 1) * inner }, checked, scratch)?;
282 Ok((0..inner).map(|i| arg.nth(if one { 0 } else { i })).collect())
283}
284
285pub fn scan_read<T: Cast, S: Cast>(
286 dims: [usize; 3], seed: Seed<'_>, direction: Direction, checked: bool, input: Read<'_>, output: &mut Buffer,
287 run: &dyn Fn(&[T], [usize; 3], Option<&[S]>, &mut Vec<S>) -> Result<(), Error>) -> Result<(), Error> {
288 scan_with_carry::<T, S, S>(dims, seed, direction, checked, input, output, run)
289}
290
291pub fn boolean_scan_read<T: Cast>(
293 dims: [usize; 3], seed: Seed<'_>, direction: Direction, checked: bool, input: Read<'_>, output: &mut Buffer,
294 run: &dyn Fn(&[T], [usize; 3], Option<&[T]>, &mut Vec<bool>) -> Result<(), Error>) -> Result<(), Error> where bool: Source<T> {
295 scan_with_carry::<T, T, bool>(dims, seed, direction, checked, input, output, run)
296}
297
298fn scan_with_carry<T: Cast, C: Cast, S: Cast + Source<C>>(
299 dims: [usize; 3], seed: Seed<'_>, direction: Direction, checked: bool, input: Read<'_>, output: &mut Buffer,
300 run: &dyn Fn(&[T], [usize; 3], Option<&[C]>, &mut Vec<S>) -> Result<(), Error>) -> Result<(), Error> {
301 use crate::agreement::{Mapping, Missing, Read};
302 let [outer, len, inner] = dims;
303 let reversed = matches!(direction, Direction::Right).then(|| {
304 input.plan.remapped(Mapping::from(vec![
305 Steps::along(outer, len * inner), Steps::Stride { start: len.saturating_sub(1) * inner, len, step: -(inner as isize) }, Steps::along(inner, 1),
306 ]), outer * len * inner, Missing::Invalid)
307 });
308 let input = Read { plan: reversed.as_ref().unwrap_or(input.plan), ..input };
309 let (mut values, mut scratch, mut seed_scratch) = (Vec::with_capacity(outer * len * inner), Vec::new(), Vec::new());
310 if inner == 0 { *output = S::buffer(values); return Ok(()); }
311 let shared = !seed.read().is_some_and(|read| read.plan.len() != 1);
313 let cells = if shared && len * inner <= CHUNK { CHUNK / (len * inner).max(1) } else { 1 };
314 for cell in (0..outer).step_by(cells) {
315 let g = cells.min(outer - cell);
316 let mut carry = seed.read().map(|read| seed_row(&read, cell, inner, checked, &mut seed_scratch)).transpose()?;
317 let rows = (CHUNK / inner).max(1);
318 for row in (0..len).step_by(rows) {
319 let n = rows.min(len - row);
320 let at = (cell * len + row) * inner;
321 match input.part::<T>(at..at + g * n * inner, checked, &mut scratch)? {
322 Arg::Many(items) => run(items, [g, n, inner], carry.as_deref(), &mut values)?,
323 Arg::One(value) => {
324 scratch.clear();
325 scratch.resize(g * n * inner, value);
326 run(&scratch, [g, n, inner], carry.as_deref(), &mut values)?;
327 }
328 }
329 let carry = carry.get_or_insert_with(Vec::new);
330 carry.clear();
331 let tail = &values[values.len() - inner..];
332 if let Some(items) = S::borrowed(tail) { carry.extend_from_slice(items); } else { carry.extend(tail.iter().copied().map(Source::<C>::read)); }
333 }
334 }
335 if reversed.is_some() { reverse_rows(&mut values, dims); }
336 *output = S::buffer(values);
337 Ok(())
338}
339
340pub unsafe fn adjacent_read<T: crate::buffer::Cast>(
345 [outer, len, inner]: [usize; 3], seed: Option<crate::agreement::Read<'_>>, checked: bool, flip: bool, input: crate::agreement::Read<'_>,
346 output: &mut crate::buffer::Buffer, block: &dyn Fn(&[T], &[T], &mut [std::mem::MaybeUninit<T>]) -> Result<(), Kind>) -> Result<(), Error> {
347 let (mut values, mut scratch, mut seed_scratch) = (Vec::with_capacity(outer * len * inner), Vec::new(), Vec::new());
348 let mut failure = None;
349 let mut pair = |earlier: &[T], items: &[T], values: &mut Vec<T>| {
350 failure = failure.or(unsafe { appended(values, items.len(), |out| if flip { block(items, earlier, out) } else { block(earlier, items, out) }) }.err());
351 };
352 for cell in (0..outer).filter(|_| len > 0) {
353 let start = seed.map(|read| seed_row(&read, cell, inner, checked, &mut seed_scratch)).transpose()?;
354 let lane: std::borrow::Cow<'_, [T]> = match input.part::<T>(cell * len * inner..(cell + 1) * len * inner, checked, &mut scratch)? {
355 Arg::Many(items) => items.into(), Arg::One(item) => vec![item; len * inner].into()
356 };
357 match &start { Some(start) => pair(start, &lane[..inner], &mut values), None => values.extend_from_slice(&lane[..inner]) }
358 pair(&lane[..(len - 1) * inner], &lane[inner..], &mut values);
359 }
360 if checked && let Some(kind) = failure { return Err(Error::new(kind, "a combination of adjacent items failed")); }
361 *output = T::buffer(values);
362 Ok(())
363}
364
365pub fn reverse_rows<T>(values: &mut [T], [outer, len, inner]: [usize; 3]) {
366 if len < 2 || inner == 0 { return; }
367 for cell in values[..outer * len * inner].chunks_exact_mut(len * inner) {
368 for row in 0..len / 2 { for column in 0..inner { cell.swap(row * inner + column, (len - row - 1) * inner + column); } }
369 }
370}