torsh_ffi/python/tensor/
storage.rs

1use crate::error::{FfiError, FfiResult};
2use crate::python::tensor::memory::MEMORY_POOL;
3use parking_lot::Mutex;
4use std::sync::atomic::{AtomicUsize, Ordering};
5use std::sync::Arc;
6use torsh_core::DType;
7
8/// Reference-counted tensor storage for memory management
9#[derive(Debug, Clone)]
10pub struct TensorStorage {
11    data: Arc<Mutex<Vec<f32>>>,
12    shape: Vec<usize>,
13    dtype: DType,
14    /// Track if this tensor was created from external memory (e.g., NumPy)
15    is_external: bool,
16    /// Gradient storage (None if no gradient computed yet)
17    grad: Arc<Mutex<Option<Vec<f32>>>>,
18    /// Version for gradient tracking (incremented on in-place operations)
19    /// Using AtomicUsize for lock-free version tracking
20    version: Arc<AtomicUsize>,
21}
22
23impl TensorStorage {
24    /// Create new tensor storage with data
25    pub fn new(data: Vec<f32>, shape: Vec<usize>, dtype: DType) -> Self {
26        Self {
27            data: Arc::new(Mutex::new(data)),
28            shape,
29            dtype,
30            is_external: false,
31            grad: Arc::new(Mutex::new(None)),
32            version: Arc::new(AtomicUsize::new(0)),
33        }
34    }
35
36    /// Create new tensor storage with pooled allocation
37    pub fn new_pooled(shape: Vec<usize>, dtype: DType) -> FfiResult<Self> {
38        let size = shape.iter().product();
39        let data = MEMORY_POOL.allocate(size)?;
40
41        Ok(Self {
42            data: Arc::new(Mutex::new(data)),
43            shape,
44            dtype,
45            is_external: false,
46            grad: Arc::new(Mutex::new(None)),
47            version: Arc::new(AtomicUsize::new(0)),
48        })
49    }
50
51    /// Create tensor storage from external data (e.g., NumPy array)
52    pub fn from_external(data: Vec<f32>, shape: Vec<usize>, dtype: DType) -> Self {
53        Self {
54            data: Arc::new(Mutex::new(data)),
55            shape,
56            dtype,
57            is_external: true,
58            grad: Arc::new(Mutex::new(None)),
59            version: Arc::new(AtomicUsize::new(0)),
60        }
61    }
62
63    /// Get a reference to the data (read-only)
64    pub fn data(&self) -> parking_lot::MutexGuard<'_, Vec<f32>> {
65        self.data.lock()
66    }
67
68    /// Get a mutable reference to the data
69    pub fn data_mut(&self) -> parking_lot::MutexGuard<'_, Vec<f32>> {
70        self.data.lock()
71    }
72
73    /// Get shape
74    pub fn shape(&self) -> &[usize] {
75        &self.shape
76    }
77
78    /// Get dtype
79    pub fn dtype(&self) -> DType {
80        self.dtype
81    }
82
83    /// Create a view with different shape (same data)
84    pub fn view(&self, new_shape: Vec<usize>) -> FfiResult<Self> {
85        let total_elements: usize = new_shape.iter().product();
86        let current_elements: usize = self.shape.iter().product();
87
88        if total_elements != current_elements {
89            return Err(FfiError::ShapeMismatch {
90                expected: vec![current_elements],
91                actual: vec![total_elements],
92            });
93        }
94
95        Ok(Self {
96            data: Arc::clone(&self.data),
97            shape: new_shape,
98            dtype: self.dtype,
99            is_external: self.is_external,
100            grad: Arc::clone(&self.grad),
101            version: Arc::clone(&self.version),
102        })
103    }
104
105    /// Get reference count
106    pub fn ref_count(&self) -> usize {
107        Arc::strong_count(&self.data)
108    }
109
110    /// Check if this storage comes from external memory
111    pub fn is_external(&self) -> bool {
112        self.is_external
113    }
114
115    /// Get gradient data (None if no gradient computed)
116    pub fn grad(&self) -> parking_lot::MutexGuard<'_, Option<Vec<f32>>> {
117        self.grad.lock()
118    }
119
120    /// Set gradient data
121    pub fn set_grad(&self, grad_data: Option<Vec<f32>>) {
122        *self.grad.lock() = grad_data;
123    }
124
125    /// Clear gradient
126    pub fn clear_grad(&self) {
127        *self.grad.lock() = None;
128    }
129
130    /// Get current version (for gradient tracking)
131    /// Uses atomic load for lock-free read access
132    pub fn version(&self) -> usize {
133        self.version.load(Ordering::Relaxed)
134    }
135
136    /// Increment version (called on in-place operations)
137    /// Uses atomic increment for lock-free update
138    pub fn increment_version(&self) {
139        self.version.fetch_add(1, Ordering::Relaxed);
140    }
141}
142
143impl Drop for TensorStorage {
144    fn drop(&mut self) {
145        // Only return to pool if this is the last reference and not external memory
146        if !self.is_external && Arc::strong_count(&self.data) == 1 {
147            if let Ok(data) = Arc::try_unwrap(std::mem::replace(
148                &mut self.data,
149                Arc::new(Mutex::new(Vec::new())),
150            )) {
151                // parking_lot::Mutex::into_inner() returns the value directly (no Result)
152                let data_vec = data.into_inner();
153                // Attempt to return to pool (ignore errors during drop)
154                let _ = MEMORY_POOL.deallocate(data_vec);
155            }
156        }
157    }
158}
torsh_ffi/python/tensor/storage.rs

torsh_ffi/python/tensor/
storage.rs
