rustymath_core/buffer/
output.rs

1use super::*;
2use crate::{Arg, parallel::Slots};
3
4pub trait OutputItem: Element { fn slots<'s, 'b>(output: &'s mut MapOutput<'_, 'b>) -> Option<&'s mut Slots<'b, Self>>; }
5
6macro_rules! outputs {
7    ([$($v:ident: ($dtype:expr) $t:ty),+]) => {
8        /// Where a map appends: an owned buffer or the next slots of a caller's result.
9        pub enum MapOutput<'a, 'b> { Buffer(&'a mut Buffer), $($v(&'a mut Slots<'b, $t>)),+ }
10        impl<'b> MapOutput<'_, 'b> {
11            fn dtype(&self) -> DType { match self { Self::Buffer(b) => b.dtype(), $(Self::$v(_) => <$t>::DTYPE),+ } }
12            pub fn reborrow(&mut self) -> MapOutput<'_, 'b> { match self { Self::Buffer(b) => MapOutput::Buffer(b), $(Self::$v(s) => MapOutput::$v(s)),+ } }
13            pub fn fill(&mut self, n: usize) {
14                match self {
15                    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); }),
16                    $(Self::$v(s) => { s.pairs(Arg::One(<$t>::FILL), Arg::One(()), n, |x, ()| (x, true)); }),+
17                }
18            }
19        }
20        $(impl OutputItem for $t {
21            fn slots<'s, 'b>(output: &'s mut MapOutput<'_, 'b>) -> Option<&'s mut Slots<'b, Self>> {
22                match output { MapOutput::$v(s) => Some(s), _ => None }
23            }
24        }
25        impl<'a, 'b> From<&'a mut Slots<'b, $t>> for MapOutput<'a, 'b> {
26            fn from(slots: &'a mut Slots<'b, $t>) -> Self { MapOutput::$v(slots) }
27        })+
28    };
29}
30element_types!((outputs));
31
32/// Conversion buffers retained for one execution, by computing dtype and operand position.
33#[derive(Default)]
34pub struct Scratch { buffers: Vec<(DType, Vec<Buffer>)>, results: Vec<Buffer>, copied: usize }
35
36impl Scratch {
37    fn buffers(&mut self, dtype: DType, n: usize) -> &mut [Buffer] {
38        let at = self.buffers.iter().position(|(stored, _)| *stored == dtype).unwrap_or_else(|| {
39            self.buffers.push((dtype, Vec::new()));
40            self.buffers.len() - 1
41        });
42        let buffers = &mut self.buffers[at].1;
43        buffers.resize_with(n.max(buffers.len()), || Buffer::with_capacity(dtype, 0));
44        &mut buffers[..n]
45    }
46    pub fn copied(&self) -> usize { self.copied }
47    pub fn record_copies(&mut self, items: usize) { self.copied += items; }
48    fn result<T: Element>(&mut self) -> &mut Vec<T> {
49        let at = self.results.iter().position(|buffer| buffer.dtype() == T::DTYPE).unwrap_or_else(|| {
50            self.results.push(Buffer::with_capacity(T::DTYPE, 0));
51            self.results.len() - 1
52        });
53        T::vec_mut(&mut self.results[at]).expect("a result buffer has its computing dtype")
54    }
55    pub fn buffer(&mut self, dtype: DType) -> &mut Buffer { &mut self.buffers(dtype, 1)[0] }
56    pub fn inputs<T: Element, const N: usize>(&mut self) -> [&mut Vec<T>; N] {
57        let buffers: &mut [Buffer; N] = self.buffers(T::DTYPE, N).try_into().expect("scratch has an entry per operand");
58        buffers.each_mut().map(|buffer| T::vec_mut(buffer).expect("scratch keeps its computing dtype"))
59    }
60}
61
62pub struct Output<'a, 'b> {
63    pub target: MapOutput<'a, 'b>, pub result_dtype: DType, pub checked: bool, pub scratch: Option<&'a mut Scratch>
64}
65
66impl Output<'_, '_> {
67    /// The block body is erased before output conversion.
68    pub fn write<O: OutputItem>(
69        &mut self, n: usize,
70        apply: &(dyn Fn(std::ops::Range<usize>, &mut Slots<'_, O>, Option<&mut Scratch>) -> bool + Sync)
71    ) -> Result<(), Error> {
72        let checked = self.checked;
73        if let MapOutput::Buffer(buffer) = &mut self.target {
74            if let Some(values) = O::vec_mut(buffer) { extended(values, n, self.scratch.as_deref_mut(), apply); } else {
75                with_dtype!(buffer.dtype(), T => {
76                    let values = T::vec_mut(buffer).expect("a buffer has its dtype");
77                    extended(values, n, self.scratch.as_deref_mut(), &|part, slots, scratch| {
78                        converted::<O>(&mut slots.into(), part, checked, scratch, apply).map_err(crate::error::record).is_ok()
79                    });
80                });
81            }
82            return Ok(());
83        }
84        if let Some(slots) = O::slots(&mut self.target) { apply(0..n, slots, self.scratch.as_deref_mut()); return Ok(()); }
85        converted(&mut self.target, 0..n, checked, self.scratch.as_deref_mut(), apply)
86    }
87}
88
89fn extended<O: Element>(
90    values: &mut Vec<O>, n: usize, scratch: Option<&mut Scratch>,
91    apply: &(dyn Fn(std::ops::Range<usize>, &mut Slots<'_, O>, Option<&mut Scratch>) -> bool + Sync)
92) {
93    match scratch {
94        Some(scratch) => { Slots::append(values, n, |slots| apply(0..n, slots, Some(scratch))); }
95        None => { crate::parallel::extended(values, n, crate::CHUNK, &|part, slots| apply(part, slots, None)); }
96    }
97}
98
99fn converted<O: OutputItem>(
100    target: &mut MapOutput<'_, '_>, range: std::ops::Range<usize>, checked: bool, mut work: Option<&mut Scratch>,
101    apply: &(dyn Fn(std::ops::Range<usize>, &mut Slots<'_, O>, Option<&mut Scratch>) -> bool + Sync)
102) -> Result<(), Error> {
103    let mut values = work.as_deref_mut().map_or_else(Vec::new, |work| std::mem::take(work.result::<O>()));
104    let mut local = Scratch::default();
105    let conversion = crate::convert::Conversion { input: O::DTYPE, output: target.dtype(), kind: crate::convert::ConversionKind::Truncate };
106    let result = (|| {
107        for at in (range.start..range.end).step_by(crate::CHUNK) {
108            let n = crate::CHUNK.min(range.end - at);
109            values.clear();
110            Slots::append(&mut values, n, |slots| apply(at..at + n, slots, work.as_deref_mut()));
111            if let Err(error) = conversion.append(O::view(&values), checked, crate::tolerant::Tolerance::EXACT,
112                work.as_deref_mut().unwrap_or(&mut local), target.reborrow()) {
113                target.fill(range.end - at - n);
114                return Err(error);
115            }
116        }
117        Ok(())
118    })();
119    if let Some(work) = work { *work.result::<O>() = values; }
120    result
121}
rustymath_core/buffer/output.rs

rustymath_core/buffer/
output.rs
