rten_simd/
span.rs

1//! Slice-like types used as inputs and outputs for vectorized operations.
2
3use std::mem::{MaybeUninit, transmute};
4
5enum SrcDestInner<'src, 'dst, T> {
6    InOut(&'src [T], &'dst mut [MaybeUninit<T>]),
7    InMut(&'dst mut [T]),
8}
9
10/// Input-output buffer for vectorized operations.
11///
12/// This can either be a single mutable buffer for operations that execute
13/// in-place (`&mut [T]`) or a pair of input and output buffers where the
14/// output is uninitialized (`([T], &mut [MaybeUninit<T>])`) and both buffers
15/// must have the same length.
16pub struct SrcDest<'src, 'dst, T: Copy> {
17    inner: SrcDestInner<'src, 'dst, T>,
18}
19
20impl<'dst, T: Copy> SrcDest<'_, 'dst, T> {
21    /// Return the source slice.
22    pub fn src(&self) -> &[T] {
23        match &self.inner {
24            SrcDestInner::InOut(src, _dest) => src,
25            SrcDestInner::InMut(src_mut) => src_mut,
26        }
27    }
28
29    /// Return the length of the input and output slices.
30    pub fn len(&self) -> usize {
31        self.src().len()
32    }
33
34    /// Return true if the input and output slices are empty.
35    pub fn is_empty(&self) -> bool {
36        self.src().is_empty()
37    }
38
39    /// Return source and destination slice pointers and the length.
40    ///
41    /// The source and destination will either alias, or the destination will
42    /// be a non-aliasing, uninitialized slice.
43    pub fn src_dest_ptr(&mut self) -> (*const T, *mut MaybeUninit<T>, usize) {
44        match &mut self.inner {
45            SrcDestInner::InOut(src, dest) => (src.as_ptr(), dest.as_mut_ptr(), src.len()),
46            SrcDestInner::InMut(src) => (
47                src.as_ptr(),
48                src.as_mut_ptr() as *mut MaybeUninit<T>,
49                src.len(),
50            ),
51        }
52    }
53
54    /// Return the initialized destination slice.
55    ///
56    /// # Safety
57    ///
58    /// If this instance was constructed with an uninitialized destination
59    /// buffer, all elements must have been initialized before this is called.
60    pub unsafe fn dest_assume_init(self) -> &'dst mut [T] {
61        match self.inner {
62            SrcDestInner::InOut(_src, dest) => unsafe {
63                transmute::<&mut [MaybeUninit<T>], &mut [T]>(dest)
64            },
65            SrcDestInner::InMut(src) => src,
66        }
67    }
68}
69
70impl<'src, 'dst, T: Copy> From<(&'src [T], &'dst mut [MaybeUninit<T>])> for SrcDest<'src, 'dst, T> {
71    #[track_caller]
72    fn from(val: (&'src [T], &'dst mut [MaybeUninit<T>])) -> Self {
73        let (src, dest) = val;
74        assert!(
75            src.len() == dest.len(),
76            "src len {} != dest len {}",
77            src.len(),
78            dest.len(),
79        );
80        SrcDest {
81            inner: SrcDestInner::InOut(src, dest),
82        }
83    }
84}
85
86impl<'dst, T: Copy> From<&'dst mut [T]> for SrcDest<'dst, 'dst, T> {
87    fn from(val: &'dst mut [T]) -> Self {
88        SrcDest {
89            inner: SrcDestInner::InMut(val),
90        }
91    }
92}
rten_simd/span.rs

rten_simd/
span.rs
