1use 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
22pub trait CowFormat: SparseFormat {
26 type Owned<T>: SparseShape + Send + Sync
28 where
29 T: Send + Sync;
30
31 fn as_storage<T: Send + Sync>(owned: &Self::Owned<T>) -> Self::Storage<'_, T>;
33
34 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
155pub enum SparseCow<'a, T: Send + Sync, F: CowFormat, Arch: SimdArch> {
184 Borrowed(SparseView<'a, T, F, Arch>),
186 Owned(F::Owned<T>),
188}
189
190impl<'a, T: Send + Sync, F: CowFormat, Arch: SimdArch> SparseCow<'a, T, F, Arch> {
191 #[inline(always)]
193 pub fn borrowed(data: F::Storage<'a, T>) -> Self {
194 Self::Borrowed(SparseView::new(data))
195 }
196
197 #[inline(always)]
199 pub fn owned(storage: F::Owned<T>) -> Self {
200 Self::Owned(storage)
201 }
202
203 #[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 #[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 #[inline(always)]
223 pub fn is_borrowed(&self) -> bool {
224 matches!(self, Self::Borrowed(_))
225 }
226
227 #[inline(always)]
229 pub fn is_owned(&self) -> bool {
230 matches!(self, Self::Owned(_))
231 }
232
233 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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}