Skip to main content

hermes_simd_core/sparse/cow/
mod.rs

1//! Clone-on-Write sparse matrix containers.
2//!
3//! One generic [`SparseCow`] covers every sparse format: the [`CowFormat`]
4//! trait maps a format marker to its heap-owned storage and the two
5//! conversions (re-borrow owned as a view, clone a view into owned). Adding a
6//! new sparse format requires one `CowFormat` impl — the container, its
7//! accessors, and the `SparseSpMv` / `SparseOps` forwarding are inherited.
8
9use super::{
10    ops::SparseOps,
11    spmv::SparseSpMv,
12    types::{BlockedCooData, CsrData, DenseWithMaskData, SellPData, SparseShape, SparseValidate},
13    BlockedCoo, Csr, DenseWithMask, SellP, SparseFormat, SparseView, Validated, ValidatedData,
14};
15use crate::arch::SimdArch;
16use crate::scalar::Scalar;
17use crate::vec::AlignedVec;
18
19pub mod owned;
20pub use owned::{OwnedBlockedCoo, OwnedCsr, OwnedDenseWithMask, OwnedSellP};
21
22/// Maps a sparse format marker to its owned storage and Cow conversions.
23///
24/// Sealed transitively through [`SparseFormat`].
25pub trait CowFormat: SparseFormat {
26    /// Heap-owned storage for this format.
27    type Owned<T>: SparseShape + Send + Sync
28    where
29        T: Send + Sync;
30
31    /// Re-borrow owned storage as the format's view storage (zero-copy).
32    fn as_storage<T: Send + Sync>(owned: &Self::Owned<T>) -> Self::Storage<'_, T>;
33
34    /// Clone a borrowed view storage into owned storage (one allocation set).
35    fn to_owned_storage<T: Clone + Send + Sync>(storage: &Self::Storage<'_, T>) -> Self::Owned<T>;
36}
37
38impl CowFormat for Csr {
39    type Owned<T>
40        = OwnedCsr<T>
41    where
42        T: Send + Sync;
43
44    #[inline(always)]
45    fn as_storage<T: Send + Sync>(owned: &Self::Owned<T>) -> CsrData<'_, T> {
46        owned.as_view()
47    }
48
49    #[inline]
50    fn to_owned_storage<T: Clone + Send + Sync>(d: &CsrData<'_, T>) -> OwnedCsr<T> {
51        OwnedCsr::new(
52            AlignedVec::from_slice_clone(d.values),
53            AlignedVec::from_slice(d.col_indices),
54            AlignedVec::from_slice(d.row_ptr),
55            d.nrows,
56            d.ncols,
57        )
58    }
59}
60
61impl<const C: usize> CowFormat for SellP<C> {
62    type Owned<T>
63        = OwnedSellP<T, C>
64    where
65        T: Send + Sync;
66
67    #[inline(always)]
68    fn as_storage<T: Send + Sync>(owned: &Self::Owned<T>) -> SellPData<'_, T, C> {
69        owned.as_view()
70    }
71
72    #[inline]
73    fn to_owned_storage<T: Clone + Send + Sync>(d: &SellPData<'_, T, C>) -> OwnedSellP<T, C> {
74        OwnedSellP::new(
75            AlignedVec::from_slice_clone(d.values),
76            AlignedVec::from_slice(d.col_indices),
77            AlignedVec::from_slice(d.slice_ptr),
78            AlignedVec::from_slice(d.slice_col_count),
79            d.nrows,
80            d.ncols,
81        )
82    }
83}
84
85impl<const BM: usize, const BN: usize> CowFormat for BlockedCoo<BM, BN> {
86    type Owned<T>
87        = OwnedBlockedCoo<T, BM, BN>
88    where
89        T: Send + Sync;
90
91    #[inline(always)]
92    fn as_storage<T: Send + Sync>(owned: &Self::Owned<T>) -> BlockedCooData<'_, T, BM, BN> {
93        owned.as_view()
94    }
95
96    #[inline]
97    fn to_owned_storage<T: Clone + Send + Sync>(
98        d: &BlockedCooData<'_, T, BM, BN>,
99    ) -> OwnedBlockedCoo<T, BM, BN> {
100        OwnedBlockedCoo::new(
101            AlignedVec::from_slice_clone(d.blocks),
102            AlignedVec::from_slice(d.block_row),
103            AlignedVec::from_slice(d.block_col),
104            d.nblocks,
105            d.nrows,
106            d.ncols,
107        )
108    }
109}
110
111impl CowFormat for DenseWithMask {
112    type Owned<T>
113        = OwnedDenseWithMask<T>
114    where
115        T: Send + Sync;
116
117    #[inline(always)]
118    fn as_storage<T: Send + Sync>(owned: &Self::Owned<T>) -> DenseWithMaskData<'_, T> {
119        owned.as_view()
120    }
121
122    #[inline]
123    fn to_owned_storage<T: Clone + Send + Sync>(
124        d: &DenseWithMaskData<'_, T>,
125    ) -> OwnedDenseWithMask<T> {
126        OwnedDenseWithMask::new(
127            AlignedVec::from_slice_clone(d.values),
128            AlignedVec::from_slice(d.mask),
129            d.nrows,
130            d.ncols,
131        )
132    }
133}
134
135impl<F> CowFormat for Validated<F>
136where
137    F: CowFormat,
138{
139    type Owned<T>
140        = ValidatedData<F::Owned<T>>
141    where
142        T: Send + Sync;
143
144    #[inline(always)]
145    fn as_storage<T: Send + Sync>(owned: &Self::Owned<T>) -> Self::Storage<'_, T> {
146        ValidatedData::new_unchecked(F::as_storage(owned.storage()))
147    }
148
149    #[inline]
150    fn to_owned_storage<T: Clone + Send + Sync>(storage: &Self::Storage<'_, T>) -> Self::Owned<T> {
151        ValidatedData::new_unchecked(F::to_owned_storage(storage.storage()))
152    }
153}
154
155/// Clone-on-Write sparse matrix, generic over the storage format `F`.
156///
157/// `Borrowed` wraps a zero-copy [`SparseView`]; `Owned` holds heap-backed
158/// storage. Reads never allocate; [`SparseCow::to_owned`] promotes exactly
159/// once.
160///
161/// # Examples
162///
163/// ```
164/// use hermes_simd_core::{Csr, CsrData, Validated};
165/// use hermes_simd_core::sparse::{SparseCow, SparseSpMv};
166/// use hermes_simd_intrinsics::Scalar;
167///
168/// let values = [1.0_f64, 2.0, 3.0];
169/// let col_indices = [0, 2, 1];
170/// let row_ptr = [0, 2, 3];
171/// let data = CsrData::new(&values, &col_indices, &row_ptr, 2, 3);
172/// let matrix = SparseCow::<f64, Validated<Csr>, Scalar>::try_borrowed(data).unwrap();
173///
174/// let x = [10.0, 20.0, 30.0];
175/// let mut y = [0.0; 2];
176/// matrix.spmv(&x, &mut y);
177///
178/// assert_eq!(matrix.nrows(), 2);
179/// assert_eq!(matrix.ncols(), 3);
180/// assert!(matrix.is_borrowed());
181/// assert_eq!(y, [70.0, 60.0]);
182/// ```
183pub enum SparseCow<'a, T: Send + Sync, F: CowFormat, Arch: SimdArch> {
184    /// Zero-copy borrowed view.
185    Borrowed(SparseView<'a, T, F, Arch>),
186    /// Owned heap-backed storage.
187    Owned(F::Owned<T>),
188}
189
190impl<'a, T: Send + Sync, F: CowFormat, Arch: SimdArch> SparseCow<'a, T, F, Arch> {
191    /// Wrap borrowed view storage (zero allocation).
192    #[inline(always)]
193    pub fn borrowed(data: F::Storage<'a, T>) -> Self {
194        Self::Borrowed(SparseView::new(data))
195    }
196
197    /// Wrap owned storage (already heap-allocated).
198    #[inline(always)]
199    pub fn owned(storage: F::Owned<T>) -> Self {
200        Self::Owned(storage)
201    }
202
203    /// Number of rows.
204    #[inline(always)]
205    pub fn nrows(&self) -> usize {
206        match self {
207            Self::Borrowed(v) => v.nrows(),
208            Self::Owned(o) => o.nrows(),
209        }
210    }
211
212    /// Number of columns.
213    #[inline(always)]
214    pub fn ncols(&self) -> usize {
215        match self {
216            Self::Borrowed(v) => v.ncols(),
217            Self::Owned(o) => o.ncols(),
218        }
219    }
220
221    /// Returns `true` if this holds a borrowed view (no heap allocation).
222    #[inline(always)]
223    pub fn is_borrowed(&self) -> bool {
224        matches!(self, Self::Borrowed(_))
225    }
226
227    /// Returns `true` if this holds heap-owned data.
228    #[inline(always)]
229    pub fn is_owned(&self) -> bool {
230        matches!(self, Self::Owned(_))
231    }
232
233    /// Promote to owned, cloning from the borrowed view if needed.
234    ///
235    /// Idempotent: an already-owned container is left untouched.
236    #[inline]
237    pub fn to_owned(&mut self)
238    where
239        T: Clone,
240    {
241        if let Self::Borrowed(v) = self {
242            *self = Self::Owned(F::to_owned_storage(v.storage()));
243        }
244    }
245}
246
247impl<'a, T: Send + Sync, F: CowFormat, Arch: SimdArch> SparseCow<'a, T, Validated<F>, Arch>
248where
249    F::Storage<'a, T>: SparseValidate,
250{
251    /// Validate borrowed storage and wrap it in a zero-copy sparse Cow.
252    ///
253    /// # Errors
254    /// Returns the format-specific validation error if `data` is malformed.
255    #[inline]
256    pub fn try_borrowed(data: F::Storage<'a, T>) -> Result<Self, crate::SimdError> {
257        Ok(Self::Borrowed(SparseView::new(ValidatedData::new(data)?)))
258    }
259}
260
261impl<'a, T, F, Arch> SparseSpMv<T> for SparseCow<'a, T, F, Arch>
262where
263    T: Scalar,
264    F: CowFormat,
265    Arch: SimdArch,
266    for<'b> SparseView<'b, T, F, Arch>: SparseSpMv<T>,
267{
268    #[inline]
269    fn spmv(&self, x: &[T], y: &mut [T]) {
270        match self {
271            Self::Borrowed(v) => v.spmv(x, y),
272            Self::Owned(o) => SparseView::<T, F, Arch>::new(F::as_storage(o)).spmv(x, y),
273        }
274    }
275}
276
277impl<'a, T, F, Arch> SparseOps<T> for SparseCow<'a, T, F, Arch>
278where
279    T: Scalar,
280    F: CowFormat,
281    Arch: SimdArch,
282    for<'b> SparseView<'b, T, F, Arch>: SparseOps<T>,
283{
284    #[inline]
285    fn sum_values(&self) -> T {
286        match self {
287            Self::Borrowed(v) => v.sum_values(),
288            Self::Owned(o) => SparseView::<T, F, Arch>::new(F::as_storage(o)).sum_values(),
289        }
290    }
291
292    #[inline]
293    fn elementwise_mul_dense(&self, dense: &[T], out_values: &mut [T]) {
294        match self {
295            Self::Borrowed(v) => v.elementwise_mul_dense(dense, out_values),
296            Self::Owned(o) => SparseView::<T, F, Arch>::new(F::as_storage(o))
297                .elementwise_mul_dense(dense, out_values),
298        }
299    }
300}
301
302impl<'a, T: Send + Sync + Clone, Arch: SimdArch> SparseCow<'a, T, Csr, Arch> {
303    /// Build an owned CSR Cow from slices.
304    #[inline]
305    pub fn from_slices(
306        values: &[T],
307        col_indices: &[i32],
308        row_ptr: &[i32],
309        nrows: usize,
310        ncols: usize,
311    ) -> Self {
312        Self::Owned(OwnedCsr::new(
313            AlignedVec::from_slice_clone(values),
314            AlignedVec::from_slice(col_indices),
315            AlignedVec::from_slice(row_ptr),
316            nrows,
317            ncols,
318        ))
319    }
320}
321
322impl<'a, T: Send + Sync + Clone, Arch: SimdArch> SparseCow<'a, T, Validated<Csr>, Arch> {
323    /// Build an owned validated CSR Cow from slices.
324    ///
325    /// # Errors
326    /// Returns the CSR validation error if the sparse structure is malformed.
327    #[inline]
328    pub fn from_slices(
329        values: &[T],
330        col_indices: &[i32],
331        row_ptr: &[i32],
332        nrows: usize,
333        ncols: usize,
334    ) -> Result<Self, crate::SimdError> {
335        Ok(Self::Owned(ValidatedData::new(OwnedCsr::new(
336            AlignedVec::from_slice_clone(values),
337            AlignedVec::from_slice(col_indices),
338            AlignedVec::from_slice(row_ptr),
339            nrows,
340            ncols,
341        ))?))
342    }
343}
344
345impl<'a, T: Send + Sync + Clone, const C: usize, Arch: SimdArch> SparseCow<'a, T, SellP<C>, Arch> {
346    /// Build an owned SELL-p Cow from slices.
347    #[inline]
348    pub fn from_slices(
349        values: &[T],
350        col_indices: &[i32],
351        slice_ptr: &[i32],
352        slice_col_count: &[i32],
353        nrows: usize,
354        ncols: usize,
355    ) -> Self {
356        Self::Owned(OwnedSellP::new(
357            AlignedVec::from_slice_clone(values),
358            AlignedVec::from_slice(col_indices),
359            AlignedVec::from_slice(slice_ptr),
360            AlignedVec::from_slice(slice_col_count),
361            nrows,
362            ncols,
363        ))
364    }
365}
366
367impl<'a, T: Send + Sync + Clone, const C: usize, Arch: SimdArch>
368    SparseCow<'a, T, Validated<SellP<C>>, Arch>
369{
370    /// Build an owned validated SELL-p Cow from slices.
371    ///
372    /// # Errors
373    /// Returns the SELL-p validation error if the sparse structure is malformed.
374    #[inline]
375    pub fn from_slices(
376        values: &[T],
377        col_indices: &[i32],
378        slice_ptr: &[i32],
379        slice_col_count: &[i32],
380        nrows: usize,
381        ncols: usize,
382    ) -> Result<Self, crate::SimdError> {
383        Ok(Self::Owned(ValidatedData::new(OwnedSellP::new(
384            AlignedVec::from_slice_clone(values),
385            AlignedVec::from_slice(col_indices),
386            AlignedVec::from_slice(slice_ptr),
387            AlignedVec::from_slice(slice_col_count),
388            nrows,
389            ncols,
390        ))?))
391    }
392}
393
394impl<'a, T: Send + Sync + Clone, const BM: usize, const BN: usize, Arch: SimdArch>
395    SparseCow<'a, T, BlockedCoo<BM, BN>, Arch>
396{
397    /// Build an owned Blocked-COO Cow from slices.
398    #[inline]
399    pub fn from_slices(
400        blocks: &[T],
401        block_row: &[i32],
402        block_col: &[i32],
403        nblocks: usize,
404        nrows: usize,
405        ncols: usize,
406    ) -> Self {
407        Self::Owned(OwnedBlockedCoo::new(
408            AlignedVec::from_slice_clone(blocks),
409            AlignedVec::from_slice(block_row),
410            AlignedVec::from_slice(block_col),
411            nblocks,
412            nrows,
413            ncols,
414        ))
415    }
416}
417
418impl<'a, T: Send + Sync + Clone, const BM: usize, const BN: usize, Arch: SimdArch>
419    SparseCow<'a, T, Validated<BlockedCoo<BM, BN>>, Arch>
420{
421    /// Build an owned validated Blocked-COO Cow from slices.
422    ///
423    /// # Errors
424    /// Returns the Blocked-COO validation error if the sparse structure is malformed.
425    #[inline]
426    pub fn from_slices(
427        blocks: &[T],
428        block_row: &[i32],
429        block_col: &[i32],
430        nblocks: usize,
431        nrows: usize,
432        ncols: usize,
433    ) -> Result<Self, crate::SimdError> {
434        Ok(Self::Owned(ValidatedData::new(OwnedBlockedCoo::new(
435            AlignedVec::from_slice_clone(blocks),
436            AlignedVec::from_slice(block_row),
437            AlignedVec::from_slice(block_col),
438            nblocks,
439            nrows,
440            ncols,
441        ))?))
442    }
443}
444
445impl<'a, T: Send + Sync + Clone, Arch: SimdArch> SparseCow<'a, T, DenseWithMask, Arch> {
446    /// Build an owned DenseWithMask Cow from slices.
447    #[inline]
448    pub fn from_slices(values: &[T], mask: &[bool], nrows: usize, ncols: usize) -> Self {
449        Self::Owned(OwnedDenseWithMask::new(
450            AlignedVec::from_slice_clone(values),
451            AlignedVec::from_slice(mask),
452            nrows,
453            ncols,
454        ))
455    }
456}
457
458impl<'a, T: Send + Sync, F: CowFormat, Arch: SimdArch> crate::sparse::types::SparseValidate
459    for SparseCow<'a, T, F, Arch>
460where
461    for<'b> F::Storage<'b, T>: crate::sparse::types::SparseValidate,
462    F::Owned<T>: crate::sparse::types::SparseValidate,
463{
464    #[inline]
465    fn validate(&self) -> Result<(), crate::SimdError> {
466        match self {
467            Self::Borrowed(v) => v.validate(),
468            Self::Owned(o) => o.validate(),
469        }
470    }
471}