hermes_simd_core/tensor/
cow.rs

1//! Clone-on-Write tensor container backed by [`AlignedVec`].
2
3use crate::align::{Alignment, Unaligned};
4use crate::vec::AlignedVec;
5
6use super::error::TensorError;
7use super::layout::{Layout, RowMajor};
8use super::strides::row_major_strides;
9use super::view::TensorView;
10
11// ---------------------------------------------------------------------------
12// TensorCow: Clone-on-Write tensor container
13// ---------------------------------------------------------------------------
14
15/// A Clone-on-Write (CoW) container for strided tensors.
16pub enum TensorCow<'a, T: 'a, const N: usize, L = RowMajor, Align: Alignment = Unaligned> {
17    /// Borrowed read-only tensor view.
18    Borrowed(TensorView<'a, T, N, L, &'a [T]>),
19    /// Owned aligned tensor buffer.
20    Owned {
21        /// Underlying aligned memory.
22        data: AlignedVec<T, Align>,
23        /// Logical shape of the tensor.
24        shape: [usize; N],
25        /// Dimension strides.
26        strides: [usize; N],
27    },
28}
29
30impl<'a, T: Copy + 'a, const N: usize, L: Layout, Align> TensorCow<'a, T, N, L, Align>
31where
32    Align: Alignment,
33{
34    /// Create a borrowed `TensorCow` wrapping a `TensorView`.
35    #[inline]
36    pub fn borrowed(view: TensorView<'a, T, N, L, &'a [T]>) -> Self {
37        Self::Borrowed(view)
38    }
39
40    /// Create an owned `TensorCow` from an `AlignedVec` and shape.
41    #[inline]
42    pub fn owned(data: AlignedVec<T, Align>, shape: [usize; N]) -> Self {
43        let strides = row_major_strides(shape);
44        Self::Owned {
45            data,
46            shape,
47            strides,
48        }
49    }
50
51    /// Create an owned `TensorCow` with explicit strides.
52    #[inline]
53    pub fn owned_with_strides(
54        data: AlignedVec<T, Align>,
55        shape: [usize; N],
56        strides: [usize; N],
57    ) -> Self {
58        Self::Owned {
59            data,
60            shape,
61            strides,
62        }
63    }
64
65    /// Obtain a read-only view of this tensor.
66    #[inline]
67    pub fn as_view(&self) -> TensorView<'_, T, N, L, &'_ [T]> {
68        match self {
69            Self::Borrowed(view) => *view,
70            Self::Owned {
71                data,
72                shape,
73                strides,
74            } => TensorView::with_strides(data.as_slice(), *shape, *strides)
75                .expect("Owned variant stores pre-validated shape and strides"),
76        }
77    }
78
79    /// Returns the total logical element count.
80    #[inline]
81    pub fn len(&self) -> usize {
82        match self {
83            Self::Borrowed(view) => view.num_elements(),
84            Self::Owned { shape, .. } => shape.iter().product(),
85        }
86    }
87
88    /// Returns true if empty.
89    #[inline]
90    pub fn is_empty(&self) -> bool {
91        self.len() == 0
92    }
93
94    /// Returns logical shape.
95    #[inline]
96    pub fn shape(&self) -> [usize; N] {
97        match self {
98            Self::Borrowed(view) => view.shape(),
99            Self::Owned { shape, .. } => *shape,
100        }
101    }
102
103    /// Returns tensor strides.
104    #[inline]
105    pub fn strides(&self) -> [usize; N] {
106        match self {
107            Self::Borrowed(view) => view.strides(),
108            Self::Owned { strides, .. } => *strides,
109        }
110    }
111
112    /// Returns whether tensor is contiguous row-major.
113    #[inline]
114    pub fn is_contiguous(&self) -> bool {
115        match self {
116            Self::Borrowed(view) => view.is_contiguous(),
117            Self::Owned { shape, strides, .. } => {
118                let expected = row_major_strides(*shape);
119                *strides == expected
120            }
121        }
122    }
123
124    /// Upgrades to `Owned` if currently borrowed and returns a mutable reference to the `AlignedVec`.
125    #[inline]
126    pub fn to_mut(&mut self) -> &mut AlignedVec<T, Align> {
127        if let Self::Borrowed(view) = *self {
128            let owned = AlignedVec::from_slice(view.as_slice());
129            *self = Self::Owned {
130                data: owned,
131                shape: view.shape(),
132                strides: view.strides(),
133            };
134        }
135        match self {
136            Self::Owned { data, .. } => data,
137            _ => unreachable!(),
138        }
139    }
140
141    /// Converts into the owned `AlignedVec` storage.
142    #[inline]
143    pub fn into_owned(self) -> AlignedVec<T, Align> {
144        match self {
145            Self::Borrowed(view) => AlignedVec::from_slice(view.as_slice()),
146            Self::Owned { data, .. } => data,
147        }
148    }
149
150    /// Reshapes the tensor to a different rank `M` without allocation.
151    #[inline]
152    pub fn reshape<const M: usize>(
153        self,
154        new_shape: [usize; M],
155    ) -> Result<TensorCow<'a, T, M, RowMajor, Align>, TensorError> {
156        if !self.is_contiguous() {
157            return Err(TensorError::NotContiguous);
158        }
159        let old_count = self.len();
160        let new_count = new_shape.iter().product::<usize>();
161        if old_count != new_count {
162            return Err(TensorError::ShapeMismatch);
163        }
164        match self {
165            Self::Borrowed(view) => {
166                let reshaped = view.reshape(new_shape)?;
167                Ok(TensorCow::Borrowed(reshaped))
168            }
169            Self::Owned { data, .. } => {
170                let strides = row_major_strides(new_shape);
171                Ok(TensorCow::Owned {
172                    data,
173                    shape: new_shape,
174                    strides,
175                })
176            }
177        }
178    }
179}
180
181// ---------------------------------------------------------------------------
182// Clone
183// ---------------------------------------------------------------------------
184
185impl<'a, T: Clone + 'a, const N: usize, L, Align> Clone for TensorCow<'a, T, N, L, Align>
186where
187    Align: Alignment,
188{
189    #[inline]
190    fn clone(&self) -> Self {
191        match self {
192            Self::Borrowed(view) => Self::Borrowed(*view),
193            Self::Owned {
194                data,
195                shape,
196                strides,
197            } => Self::Owned {
198                data: data.clone(),
199                shape: *shape,
200                strides: *strides,
201            },
202        }
203    }
204}
205
206// ---------------------------------------------------------------------------
207// Deref
208// ---------------------------------------------------------------------------
209
210impl<'a, T: 'a, const N: usize, L, Align> core::ops::Deref for TensorCow<'a, T, N, L, Align>
211where
212    Align: Alignment,
213{
214    type Target = [T];
215
216    #[inline]
217    fn deref(&self) -> &Self::Target {
218        match self {
219            Self::Borrowed(view) => view.as_slice(),
220            Self::Owned { data, .. } => data.as_slice(),
221        }
222    }
223}
224
225// ---------------------------------------------------------------------------
226// PartialEq / Eq
227// ---------------------------------------------------------------------------
228
229impl<'a, 'b, T, const N: usize, L1, L2, A1, A2> PartialEq<TensorCow<'b, T, N, L2, A2>>
230    for TensorCow<'a, T, N, L1, A1>
231where
232    T: PartialEq,
233    A1: Alignment,
234    A2: Alignment,
235{
236    #[inline]
237    fn eq(&self, other: &TensorCow<'b, T, N, L2, A2>) -> bool {
238        let s1: &[T] = self;
239        let s2: &[T] = other;
240        s1 == s2
241    }
242}
243
244impl<'a, T, const N: usize, L, Align> Eq for TensorCow<'a, T, N, L, Align>
245where
246    T: Eq,
247    Align: Alignment,
248{
249}
hermes_simd_core/tensor/cow.rs

hermes_simd_core/tensor/
cow.rs
