rten_simd/
writer.rs

1use std::mem::{MaybeUninit, transmute};
2
3use crate::Elem;
4use crate::ops::BitOps;
5
6/// Utility for incrementally filling an uninitialized slice, one SIMD vector
7/// at a time.
8pub struct SliceWriter<'a, T> {
9    buf: &'a mut [MaybeUninit<T>],
10    n_init: usize,
11}
12
13impl<'a, T: Elem> SliceWriter<'a, T> {
14    /// Create a writer which initializes elements of `buf`.
15    pub fn new(buf: &'a mut [MaybeUninit<T>]) -> Self {
16        SliceWriter { buf, n_init: 0 }
17    }
18
19    /// Initialize the next `ops.len()` elements of the slice from the contents
20    /// of SIMD vector `xs`.
21    ///
22    /// Panics if the slice does not have space for `ops.len()` elements.
23    pub fn write_vec<O: BitOps<T>>(&mut self, ops: O, xs: O::Simd) {
24        let written = ops.store_uninit(xs, &mut self.buf[self.n_init..]);
25        self.n_init += written.len();
26    }
27
28    /// Initialize the next `N * ops.len()` elements of the slice from the
29    /// contents of the SIMD vectors `xs`.
30    ///
31    /// This is equivalent to calling [`write_vec`](Self::write_vec) for each
32    /// vector, but performs a single bounds check for the whole batch, which is
33    /// useful when writing several vectors per loop iteration.
34    ///
35    /// Panics if the slice does not have space for `N * ops.len()` elements.
36    pub fn write_vecs<O: BitOps<T>, const N: usize>(&mut self, ops: O, xs: [O::Simd; N]) {
37        let written = ops.store_many_uninit(xs, &mut self.buf[self.n_init..]);
38        self.n_init += written.len();
39    }
40
41    /// Initialize the next element of the slice from `x`.
42    ///
43    /// Panics if the slice does not have space for writing any more elements.
44    pub fn write_scalar(&mut self, x: T) {
45        self.buf[self.n_init].write(x);
46        self.n_init += 1;
47    }
48
49    /// Finish writing the slice and return the initialized portion.
50    pub fn into_mut_slice(self) -> &'a mut [T] {
51        let init = &mut self.buf[0..self.n_init];
52
53        // Safety: All elements in `init` have been initialized.
54        unsafe { transmute::<&mut [MaybeUninit<T>], &mut [T]>(init) }
55    }
56}
57
58#[cfg(test)]
59mod tests {
60    use std::mem::MaybeUninit;
61
62    use crate::ops::BitOps;
63    use crate::{Isa, SimdOp, SliceWriter};
64
65    #[test]
66    fn test_slice_writer() {
67        struct MemCopy<'src, 'dest> {
68            src: &'src [f32],
69            dest: &'dest mut [MaybeUninit<f32>],
70        }
71
72        impl<'src, 'dest> SimdOp for MemCopy<'src, 'dest> {
73            type Output = &'dest mut [f32];
74
75            fn eval<I: Isa>(self, isa: I) -> &'dest mut [f32] {
76                let ops = isa.f32();
77
78                let mut src_chunks = self.src.chunks_exact(ops.len());
79                let mut dest_writer = SliceWriter::new(self.dest);
80
81                for chunk in src_chunks.by_ref() {
82                    let xs = ops.load(chunk);
83                    dest_writer.write_vec(ops, xs);
84                }
85
86                for x in src_chunks.remainder() {
87                    dest_writer.write_scalar(*x);
88                }
89
90                dest_writer.into_mut_slice()
91            }
92        }
93
94        // Length which should cover the vectorized body and tail cases for
95        // every ISA.
96        let len = 17;
97        let src: Vec<_> = (0..len).map(|x| x as f32).collect();
98        let mut dest = Vec::with_capacity(src.len());
99
100        let copied = MemCopy {
101            src: &src,
102            dest: dest.spare_capacity_mut(),
103        }
104        .dispatch();
105        assert_eq!(copied, src);
106    }
107}
rten_simd/writer.rs

rten_simd/
writer.rs
