rustymath_kernels/operation/map/
io.rs

1use super::*;
2use crate::{buffer::Cast, parallel::Slots};
3
4#[derive(Clone, Copy)]
5pub(crate) enum Inputs<'a, 'b> { Mapped(&'a [crate::agreement::Read<'b>]), Ready(&'a [View<'b>]) }
6
7impl<'a, 'b> Inputs<'a, 'b> {
8    pub(crate) fn len(self) -> usize { match self { Self::Mapped(x) => x.len(), Self::Ready(x) => x.len() } }
9    pub(crate) fn selected(self, k: usize) -> Self { match self { Self::Mapped(x) => Self::Mapped(&x[k..k + 1]), Self::Ready(x) => Self::Ready(&x[k..k + 1]) } }
10    pub(crate) fn real(self) -> bool {
11        let real = |v: View<'_>| !matches!(v.dtype(), DType::Complex(_));
12        match self { Self::Mapped(x) => x.iter().all(|r| r.sources.iter().all(|s| real(s.view))), Self::Ready(x) => x.iter().all(|&v| real(v)) }
13    }
14    pub(crate) fn in_place<F: Cast>(self, range: &std::ops::Range<usize>) -> bool {
15        match self { Self::Mapped(x) => x.iter().all(|r| r.in_place::<F>(range)), Self::Ready(x) => x.iter().all(|&v| F::of(v).is_some()) }
16    }
17    pub(crate) fn part<'s, F: Cast, const N: usize>(
18        self, range: std::ops::Range<usize>, checked: bool, scratch: &'s mut [Vec<F>; N]) -> Result<[Arg<'s, F>; N], Error> where 'a: 's, 'b: 's {
19        let mut args = [Arg::Many(&[]); N];
20        for (k, (arg, scratch)) in args.iter_mut().zip(scratch).enumerate() {
21            *arg = match self {
22                Self::Mapped(x) => x[k].part(range.clone(), checked, scratch)?,
23                Self::Ready(x) => {
24                    let view = if x[k].len() == 1 { x[k] } else { x[k].slice(range.clone()) };
25                    match F::of(view) {
26                        Some(items) => Arg::of(items),
27                        None => {
28                            scratch.clear();
29                            F::extend_cast(view, view.len(), checked, scratch)?;
30                            Arg::of(scratch)
31                        }
32                    }
33                }
34            };
35        }
36        Ok(args)
37    }
38}
39
40/// Where a map appends its results: an owned buffer or the next slots of a caller's result.
41pub enum MapOutput<'a, 'b> {
42    Buffer(&'a mut Buffer), F16(&'a mut Slots<'b, half::f16>), F32(&'a mut Slots<'b, f32>), F64(&'a mut Slots<'b, f64>), Bool(&'a mut Slots<'b, bool>),
43    U8(&'a mut Slots<'b, u8>), I16(&'a mut Slots<'b, i16>), I32(&'a mut Slots<'b, i32>), I64(&'a mut Slots<'b, i64>),
44    C32(&'a mut Slots<'b, num_complex::Complex32>), C64(&'a mut Slots<'b, num_complex::Complex64>), Char(&'a mut Slots<'b, char>)
45}
46
47pub(crate) trait OutputItem: Element { fn slots<'s, 'b>(output: &'s mut MapOutput<'_, 'b>) -> Option<&'s mut Slots<'b, Self>>; }
48
49macro_rules! outputs {
50    ($($v:ident: $t:ty),+ $(,)?) => {
51        impl MapOutput<'_, '_> {
52            fn dtype(&self) -> DType { match self { Self::Buffer(b) => b.dtype(), $(Self::$v(_) => <$t>::DTYPE),+ } }
53            fn copy(&mut self, view: View<'_>) {
54                match self { Self::Buffer(b) => { b.extend_from(view, 0..view.len()); }, $(Self::$v(s) => s.copy(<$t>::of(view).expect("result conversion has the destination's dtype"))),+ }
55            }
56            pub(crate) fn fill(&mut self, n: usize) {
57                match self {
58                    Self::Buffer(b) => with_dtype!(b.dtype(), T => { let v = T::vec_mut(b).expect("a buffer has its dtype"); v.resize(v.len() + n, T::FILL); }),
59                    $(Self::$v(s) => { s.pairs(Arg::One(<$t>::FILL), Arg::One(()), n, |x, ()| (x, true)); }),+
60                }
61            }
62        }
63        $(impl OutputItem for $t {
64            fn slots<'s, 'b>(output: &'s mut MapOutput<'_, 'b>) -> Option<&'s mut Slots<'b, Self>> {
65                match output { MapOutput::$v(s) => Some(s), _ => None }
66            }
67        }
68        impl<'a, 'b> From<&'a mut Slots<'b, $t>> for MapOutput<'a, 'b> {
69            fn from(slots: &'a mut Slots<'b, $t>) -> Self { MapOutput::$v(slots) }
70        })+
71    };
72}
73outputs!(F16: half::f16, F32: f32, F64: f64, Bool: bool, U8: u8, I16: i16, I32: i32, I64: i64,
74    C32: num_complex::Complex32, C64: num_complex::Complex64, Char: char);
75
76pub(crate) struct Output<'a, 'b> { pub target: MapOutput<'a, 'b>, pub dtype: DType, pub checked: bool }
77
78impl Output<'_, '_> {
79    /// The block body is already erased here: output conversion never instantiates a formula at a second computing type.
80    pub(crate) fn write<O: OutputItem>(&mut self, n: usize, apply: &(dyn Fn(std::ops::Range<usize>, &mut Slots<'_, O>) -> bool + Sync)) -> Result<(), Error> {
81        let checked = self.checked;
82        if let MapOutput::Buffer(buffer) = &mut self.target {
83            if let Some(values) = O::vec_mut(buffer) { crate::parallel::extended(values, n, crate::CHUNK, apply); } else {
84                with_dtype!(buffer.dtype(), T => {
85                    let values = T::vec_mut(buffer).expect("a buffer has its dtype");
86                    crate::parallel::extended(values, n, crate::CHUNK, &|part, slots| {
87                        converted::<O>(&mut slots.into(), part, checked, apply).map_err(crate::error::record).is_ok()
88                    });
89                });
90            }
91            return Ok(());
92        }
93        if let Some(slots) = O::slots(&mut self.target) { apply(0..n, slots); return Ok(()); }
94        converted(&mut self.target, 0..n, checked, apply)
95    }
96}
97
98fn converted<O: OutputItem>(
99    target: &mut MapOutput<'_, '_>, range: std::ops::Range<usize>, checked: bool, apply: &(dyn Fn(std::ops::Range<usize>, &mut Slots<'_, O>) -> bool + Sync)
100) -> Result<(), Error> {
101    let mut scratch = Vec::<O>::new();
102    for at in (range.start..range.end).step_by(crate::CHUNK) {
103        let n = crate::CHUNK.min(range.end - at);
104        scratch.clear();
105        Slots::append(&mut scratch, n, |slots| apply(at..at + n, slots));
106        match Buffer::cast_results(O::view(&scratch), target.dtype(), checked) {
107            Ok(buffer) => target.copy(buffer.view()), Err(error) => { target.fill(range.end - at); return Err(error); }
108        }
109    }
110    Ok(())
111}
rustymath_kernels/operation/map/io.rs

rustymath_kernels/operation/map/
io.rs
