ariadnetor_tensor/block_sparse/
storage.rs

1//! `BlockSparseStorage<T>`: pure-data half of the block-sparse tensor split.
2//!
3//! Carries only the packed flat element buffer. Block metadata,
4//! sector indices, flux, shape, and memory order live on the paired
5//! [`BlockSparseLayout<S>`](crate::BlockSparseLayout); the wrapper
6//! [`BlockSparseTensorData<T, S>`](crate::BlockSparseTensorData) joins
7//! the two.
8//!
9//! `T: Scalar`-bound scalar-only data operations (`stored_len`,
10//! `norm`, `norm_frobenius`, `normalize`) live here because their
11//! bodies touch only `data`; they require no layout field, no
12//! symmetry sector, and no compute backend.
13
14use std::sync::Arc;
15
16use aligned_vec::{AVec, ConstAlign};
17use num_traits::{Float, One, Zero};
18
19use crate::{Sector, Storage, StorageFor};
20
21/// Pure-data half of the block-sparse tensor split.
22///
23/// Holds a 64-byte-aligned packed buffer of allowed-block elements
24/// with Arc-based shared ownership (Copy-on-Write via
25/// [`Arc::make_mut`]). The interpretation (which slab maps to which
26/// block coordinate) lives on
27/// [`BlockSparseLayout<S>`](crate::BlockSparseLayout); the
28/// `BlockSparseTensorData<T, S>` wrapper joins them.
29///
30/// Notice that `BlockSparseStorage<T>` is sector-agnostic — the
31/// symmetry sector parameter `S` lives only on the layout, not on
32/// the storage.
33pub struct BlockSparseStorage<T> {
34    data: Arc<AVec<T, ConstAlign<64>>>,
35}
36
37// Manual Clone: Arc::clone does not require T: Clone (same pattern as DenseStorage<T>).
38impl<T> Clone for BlockSparseStorage<T> {
39    fn clone(&self) -> Self {
40        Self {
41            data: Arc::clone(&self.data),
42        }
43    }
44}
45
46impl<T> BlockSparseStorage<T> {
47    /// Construct from a `Vec<T>`, internally rebuilding into a
48    /// 64-byte-aligned buffer.
49    pub fn new(data: Vec<T>) -> Self
50    where
51        T: Clone,
52    {
53        let aligned = AVec::from_slice(64, &data);
54        Self {
55            data: Arc::new(aligned),
56        }
57    }
58
59    /// Construct from an already-aligned `AVec` (zero-copy).
60    pub(crate) fn from_aligned(data: AVec<T, ConstAlign<64>>) -> Self {
61        Self {
62            data: Arc::new(data),
63        }
64    }
65
66    /// Reference to the packed flat buffer.
67    pub fn data(&self) -> &[T] {
68        &self.data[..]
69    }
70
71    /// Mutable reference to the packed flat buffer (triggers CoW if
72    /// shared).
73    pub fn data_mut(&mut self) -> &mut [T]
74    where
75        T: Clone,
76    {
77        Arc::make_mut(&mut self.data).as_mut_slice()
78    }
79
80    /// Mutable access to the underlying `Arc` (for CoW-aware paths).
81    pub(crate) fn arc_mut(&mut self) -> &mut Arc<AVec<T, ConstAlign<64>>> {
82        &mut self.data
83    }
84}
85
86impl<T> Storage for BlockSparseStorage<T> {
87    type Element = T;
88
89    fn flat_len(&self) -> usize {
90        self.data.len()
91    }
92}
93
94impl<T, S: Sector> StorageFor<crate::BlockSparseLayout<S>> for BlockSparseStorage<T> {}
95
96// ---------------------------------------------------------------------------
97// Scalar-only data operations
98//
99// These read only `self.data` (no block metadata, no flux, no backend),
100// so they live on the storage half. The joined-form
101// `BlockSparseTensorData` re-exposes them via thin forwarders.
102// ---------------------------------------------------------------------------
103
104impl<T> BlockSparseStorage<T>
105where
106    T: Clone,
107{
108    /// Scale every stored element by a scalar factor in place
109    /// (triggers CoW if shared).
110    pub(crate) fn scale<S>(&mut self, factor: S)
111    where
112        T: std::ops::Mul<S, Output = T>,
113        S: Clone,
114    {
115        let data = Arc::make_mut(&mut self.data).as_mut_slice();
116        for elem in data.iter_mut() {
117            *elem = elem.clone() * factor.clone();
118        }
119    }
120}
121
122impl<T> BlockSparseStorage<T>
123where
124    T: ariadnetor_core::Scalar,
125{
126    /// Total number of stored elements across all blocks (= length of
127    /// the flat packed buffer).
128    pub(crate) fn stored_len(&self) -> usize {
129        self.data.len()
130    }
131
132    /// Squared Frobenius norm: Σ |element|².
133    fn norm_squared(&self) -> T::Real {
134        self.data
135            .iter()
136            .map(|&x| {
137                let a = x.abs();
138                a * a
139            })
140            .fold(T::Real::zero(), |acc, x| acc + x)
141    }
142
143    /// Frobenius norm: √(Σ |element|²).
144    pub(crate) fn norm_frobenius(&self) -> T::Real {
145        self.norm_squared().sqrt()
146    }
147
148    /// Frobenius norm (alias for [`norm_frobenius`](Self::norm_frobenius)).
149    pub(crate) fn norm(&self) -> T::Real {
150        self.norm_frobenius()
151    }
152
153    /// Normalize to unit Frobenius norm (in-place).
154    ///
155    /// Returns the norm before normalization.
156    /// Panics if the tensor has zero norm.
157    pub(crate) fn normalize(&mut self) -> T::Real {
158        let norm = self.norm_frobenius();
159        assert!(norm != T::Real::zero(), "Cannot normalize zero tensor");
160        let inv_norm = T::Real::one() / norm;
161        let data = Arc::make_mut(&mut self.data);
162        for elem in data.iter_mut() {
163            *elem = elem.scale_real(inv_norm);
164        }
165        norm
166    }
167}
ariadnetor_tensor/block_sparse/storage.rs

ariadnetor_tensor/block_sparse/
storage.rs
