tract_data/tensor/
storage.rs

1use std::alloc::Layout;
2use std::fmt;
3use std::hash::Hash;
4
5use crate::blob::Blob;
6
7/// Trait abstracting over tensor storage backends.
8///
9/// `DenseStorage` is the primary implementation backed by a contiguous `Blob`.
10/// Non-dense backends are held behind `StorageKind::Other(Box<dyn TensorStorage>)`.
11pub trait TensorStorage: Send + Sync + fmt::Debug + fmt::Display {
12    fn byte_len(&self) -> usize;
13    fn is_empty(&self) -> bool;
14    fn deep_clone(&self) -> Box<dyn TensorStorage>;
15    fn same_as(&self, other: &dyn TensorStorage) -> bool;
16    fn as_dense(&self) -> Option<&DenseStorage>;
17    fn as_dense_mut(&mut self) -> Option<&mut DenseStorage>;
18    fn into_dense(self: Box<Self>) -> Option<DenseStorage>;
19}
20
21/// Dense, contiguous storage backed by a `Blob`.
22#[derive(Eq)]
23pub struct DenseStorage(pub(crate) Blob);
24
25impl DenseStorage {
26    #[inline]
27    pub fn layout(&self) -> &Layout {
28        self.0.layout()
29    }
30
31    #[inline]
32    pub fn as_bytes(&self) -> &[u8] {
33        self.0.as_bytes()
34    }
35
36    #[inline]
37    pub fn as_bytes_mut(&mut self) -> &mut [u8] {
38        self.0.as_bytes_mut()
39    }
40
41    #[inline]
42    pub fn as_ptr(&self) -> *const u8 {
43        self.0.as_bytes().as_ptr()
44    }
45
46    #[inline]
47    pub fn as_mut_ptr(&mut self) -> *mut u8 {
48        self.0.as_bytes_mut().as_mut_ptr()
49    }
50
51    #[inline]
52    pub fn into_blob(self) -> Blob {
53        self.0
54    }
55}
56
57impl Default for DenseStorage {
58    #[inline]
59    fn default() -> Self {
60        DenseStorage(Blob::default())
61    }
62}
63
64impl Clone for DenseStorage {
65    #[inline]
66    fn clone(&self) -> Self {
67        DenseStorage(self.0.clone())
68    }
69}
70
71impl Hash for DenseStorage {
72    #[inline]
73    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
74        self.0.hash(state);
75    }
76}
77
78impl PartialEq for DenseStorage {
79    #[inline]
80    fn eq(&self, other: &Self) -> bool {
81        self.0 == other.0
82    }
83}
84
85impl From<Blob> for DenseStorage {
86    #[inline]
87    fn from(blob: Blob) -> Self {
88        DenseStorage(blob)
89    }
90}
91
92impl std::ops::Deref for DenseStorage {
93    type Target = [u8];
94    #[inline]
95    fn deref(&self) -> &[u8] {
96        self.0.as_bytes()
97    }
98}
99
100impl std::ops::DerefMut for DenseStorage {
101    #[inline]
102    fn deref_mut(&mut self) -> &mut [u8] {
103        self.0.as_bytes_mut()
104    }
105}
106
107impl fmt::Debug for DenseStorage {
108    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
109        fmt::Debug::fmt(&self.0, f)
110    }
111}
112
113impl fmt::Display for DenseStorage {
114    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
115        fmt::Display::fmt(&self.0, f)
116    }
117}
118
119impl TensorStorage for DenseStorage {
120    #[inline]
121    fn is_empty(&self) -> bool {
122        self.0.is_empty()
123    }
124
125    #[inline]
126    fn byte_len(&self) -> usize {
127        self.0.len()
128    }
129
130    fn deep_clone(&self) -> Box<dyn TensorStorage> {
131        Box::new(DenseStorage(self.0.clone()))
132    }
133
134    fn same_as(&self, other: &dyn TensorStorage) -> bool {
135        if let Some(other) = other.as_dense() {
136            !self.0.is_empty()
137                && self.0.as_bytes().as_ptr() == other.as_ptr()
138                && self.0.len() == other.byte_len()
139        } else {
140            false
141        }
142    }
143
144    fn as_dense(&self) -> Option<&DenseStorage> {
145        Some(self)
146    }
147
148    fn as_dense_mut(&mut self) -> Option<&mut DenseStorage> {
149        Some(self)
150    }
151
152    fn into_dense(self: Box<Self>) -> Option<DenseStorage> {
153        Some(*self)
154    }
155}
156
157/// Inline enum replacing `Box<dyn TensorStorage>`.
158///
159/// The common `Dense` case stays inline (no heap alloc, no vtable indirection).
160/// `Other` covers future non-dense backends behind a single Box indirection.
161#[allow(dead_code)]
162pub(crate) enum StorageKind {
163    Dense(DenseStorage),
164    Other(Box<dyn TensorStorage>),
165}
166
167impl StorageKind {
168    #[inline]
169    pub fn as_dense(&self) -> Option<&DenseStorage> {
170        match self {
171            StorageKind::Dense(d) => Some(d),
172            StorageKind::Other(o) => o.as_dense(),
173        }
174    }
175
176    #[inline]
177    pub fn as_dense_mut(&mut self) -> Option<&mut DenseStorage> {
178        match self {
179            StorageKind::Dense(d) => Some(d),
180            StorageKind::Other(o) => o.as_dense_mut(),
181        }
182    }
183
184    #[inline]
185    pub fn into_dense(self) -> Option<DenseStorage> {
186        match self {
187            StorageKind::Dense(d) => Some(d),
188            StorageKind::Other(o) => o.into_dense(),
189        }
190    }
191
192    #[inline]
193    pub fn byte_len(&self) -> usize {
194        match self {
195            StorageKind::Dense(d) => d.0.len(),
196            StorageKind::Other(o) => o.byte_len(),
197        }
198    }
199
200    #[inline]
201    pub fn is_empty(&self) -> bool {
202        match self {
203            StorageKind::Dense(d) => d.0.is_empty(),
204            StorageKind::Other(o) => o.is_empty(),
205        }
206    }
207
208    #[inline]
209    #[allow(dead_code)]
210    pub fn deep_clone(&self) -> StorageKind {
211        match self {
212            StorageKind::Dense(d) => StorageKind::Dense(d.clone()),
213            StorageKind::Other(o) => StorageKind::Other(o.deep_clone()),
214        }
215    }
216}
tract_data/tensor/storage.rs

tract_data/tensor/
storage.rs
