luma_tensor/tensor/
mod.rs

1mod dim;
2mod layout;
3mod shape;
4pub use dim::*;
5pub use layout::*;
6pub use shape::*;
7
8use std::sync::atomic::{AtomicUsize, Ordering};
9use std::sync::{Arc, RwLock};
10
11use crate::dtype::Storage;
12use crate::{Bool, Cpu, DTypeKind, Device, Float, Int};
13
14/// Unique, monotonically increasing tensor identity. Used as a map key during
15/// autograd (a tensor's storage/layout may change, but its id never does).
16#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
17pub struct TensorId(usize);
18
19impl TensorId {
20    pub(crate) fn new() -> Self {
21        static COUNTER: AtomicUsize = AtomicUsize::new(0);
22        Self(COUNTER.fetch_add(1, Ordering::Relaxed))
23    }
24
25    pub fn as_usize(&self) -> usize {
26        self.0
27    }
28}
29
30/// The reference-counted tensor handle. Cheap to clone (bumps an `Arc`).
31///
32/// Two generic parameters:
33/// - `D: Device` — where the data lives (e.g. `Cpu`); selects the storage type.
34/// - `K: DTypeKind<D>` — the *kind* (`Float`/`Int`/`Bool`), a compile-time marker.
35///   The concrete precision within a kind (`f32` vs `f64`, ...) is a runtime
36///   [`DType`] carried inside the tensor.
37///
38/// `Float` is the default kind so `Tensor<Cpu>` means a float tensor.
39pub struct Tensor<D: Device = Cpu, K: DTypeKind<D> = Float>(pub(crate) Arc<TensorImpl<D, K>>);
40pub type FloatTensor<D = Cpu> = Tensor<D, Float>;
41pub type IntTensor<D = Cpu> = Tensor<D, Int>;
42pub type BoolTensor<D = Cpu> = Tensor<D, Bool>;
43
44pub struct TensorImpl<D, K>
45where
46    D: Device,
47    K: DTypeKind<D>,
48{
49    pub(crate) id: TensorId,
50
51    pub(crate) storage: Option<Arc<RwLock<K::Storage>>>,
52    pub(crate) layout: Layout,
53
54    pub(crate) meta: K::Meta,
55
56    pub(crate) dtype: K::DType,
57    pub(crate) device: D,
58}
59
60impl<D: Device, K: DTypeKind<D>> Tensor<D, K> {
61    pub fn id(&self) -> TensorId {
62        self.0.id
63    }
64
65    pub fn dtype(&self) -> K::DType {
66        self.0.dtype
67    }
68
69    pub fn device(&self) -> &D {
70        &self.0.device
71    }
72
73    pub fn layout(&self) -> &Layout {
74        &self.0.layout
75    }
76
77    pub fn shape(&self) -> &Shape {
78        self.0.layout.shape()
79    }
80
81    pub fn stride(&self) -> &[usize] {
82        self.0.layout.stride()
83    }
84
85    pub fn rank(&self) -> usize {
86        self.shape().rank()
87    }
88
89    pub fn dims(&self) -> &[usize] {
90        self.shape().dims()
91    }
92
93    pub fn dim<Di: Dim>(&self, dim: Di) -> crate::Result<usize> {
94        self.shape().dim(dim)
95    }
96
97    pub fn element_count(&self) -> usize {
98        self.shape().element_count()
99    }
100
101    pub fn is_contiguous(&self) -> bool {
102        self.0.layout.is_contiguous()
103    }
104
105    /// A meta tensor carries shape but no storage.
106    pub fn is_meta(&self) -> bool {
107        self.0.storage.is_none()
108    }
109
110    /// Access the underlying storage, erroring on a meta tensor.
111    pub fn storage(&self) -> crate::Result<&Arc<RwLock<K::Storage>>> {
112        self.0.storage.as_ref().ok_or(crate::Error::MetaTensor)
113    }
114
115    /// Read-lock the storage, erroring on a meta tensor.
116    pub(crate) fn storage_read(&self) -> crate::Result<std::sync::RwLockReadGuard<'_, K::Storage>> {
117        Ok(self.storage()?.read().expect("storage read lock"))
118    }
119
120    /// Write-lock the storage, erroring on a meta tensor.
121    pub(crate) fn storage_write(&self) -> crate::Result<std::sync::RwLockWriteGuard<'_, K::Storage>> {
122        Ok(self.storage()?.write().expect("storage write lock"))
123    }
124
125    /// Verify `self` and `rhs` have the same shape; return a reference to it.
126    pub(crate) fn same_shape(&self, rhs: &Self, op: &'static str) -> crate::Result<&Shape> {
127        if self.shape() != rhs.shape() {
128            return Err(crate::Error::ShapeMismatchBinaryOp { lhs: self.shape().clone(), rhs: rhs.shape().clone(), op });
129        }
130        Ok(self.shape())
131    }
132
133    /// Build a tensor from a freshly-computed storage + layout + autograd meta.
134    pub(crate) fn from_storage<L: Into<Layout>>(storage: K::Storage, layout: L, meta: K::Meta) -> Self {
135        let device = storage.device().clone();
136        Tensor(Arc::new(TensorImpl {
137            id: TensorId::new(),
138            dtype: storage.dtype(),
139            storage: Some(Arc::new(RwLock::new(storage))),
140            layout: layout.into(),
141            meta,
142            device,
143        }))
144    }
145
146    pub(crate) fn phantom_storage<L: Into<Layout>>(layout: L, dtype: K::DType, device: D) -> Self {
147        Tensor(Arc::new(TensorImpl {
148            id: TensorId::new(),
149            dtype,
150            storage: None,
151            layout: layout.into(),
152            meta: K::Meta::default(),
153            device,
154        }))
155    }
156
157    /// Build a view that installs a new layout and meta (used by
158    /// transpose/slice/broadcast/etc.).
159    ///
160    /// The caller resolves `storage`: compute devices pass `self.0.storage.clone()`
161    /// (alias), while a tracing device passes a freshly-built storage so the view
162    /// is a distinct graph value. See `ShapeDTypeKind::view_dispatch`.
163    pub(crate) fn share_storage<L: Into<Layout>>(&self, layout: L, meta: K::Meta, storage: Option<Arc<RwLock<K::Storage>>>) -> Self {
164        Tensor(Arc::new(TensorImpl {
165            id: TensorId::new(),
166            dtype: self.0.dtype,
167            storage,
168            layout: layout.into(),
169            meta,
170            device: self.device().clone(),
171        }))
172    }
173}
174
175impl<D: Device, K: DTypeKind<D>> Clone for Tensor<D, K> {
176    fn clone(&self) -> Self {
177        Self(self.0.clone())
178    }
179}
180
181impl<D: Device, K: DTypeKind<D>> std::hash::Hash for Tensor<D, K> {
182    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
183        self.0.id.hash(state)
184    }
185}
186
187impl<D: Device, K: DTypeKind<D>> PartialEq for Tensor<D, K> {
188    fn eq(&self, other: &Self) -> bool {
189        self.0.id == other.0.id
190    }
191}
192
193impl<D: Device, K: DTypeKind<D>> Eq for Tensor<D, K> {}
194
195impl<D: Device, K: DTypeKind<D>> AsRef<Tensor<D, K>> for Tensor<D, K> {
196    fn as_ref(&self) -> &Tensor<D, K> {
197        self
198    }
199}
luma_tensor/tensor/mod.rs

luma_tensor/tensor/
mod.rs
