rstsr_core/tensor/
ownership_conversion.rs

1use crate::prelude_dev::*;
2
3/* #region basic conversion */
4
5/// Methods for tensor ownership conversion.
6impl<R, T, B, D> TensorAny<R, T, B, D>
7where
8    D: DimAPI,
9    B: DeviceAPI<T>,
10    R: DataAPI<Data = B::Raw>,
11{
12    /// Get a view of tensor.
13    pub fn view(&self) -> TensorView<'_, T, B, D> {
14        let layout = self.layout().clone();
15        let data = self.data().as_ref();
16        let storage = Storage::new(data, self.device().clone());
17        unsafe { TensorBase::new_unchecked(storage, layout) }
18    }
19
20    /// Get a mutable view of tensor.
21    pub fn view_mut(&mut self) -> TensorMut<'_, T, B, D>
22    where
23        R: DataMutAPI,
24    {
25        let device = self.device().clone();
26        let layout = self.layout().clone();
27        let data = self.data_mut().as_mut();
28        let storage = Storage::new(data, device);
29        unsafe { TensorBase::new_unchecked(storage, layout) }
30    }
31
32    /// Convert current tensor into copy-on-write.
33    pub fn into_cow<'a>(self) -> TensorCow<'a, T, B, D>
34    where
35        R: DataIntoCowAPI<'a>,
36    {
37        let (storage, layout) = self.into_raw_parts();
38        let (data, device) = storage.into_raw_parts();
39        let storage = Storage::new(data.into_cow(), device);
40        unsafe { TensorBase::new_unchecked(storage, layout) }
41    }
42
43    /// Convert tensor into owned tensor.
44    ///
45    /// Data is either moved or fully cloned.
46    /// Layout is not involved; i.e. all underlying data is moved or cloned
47    /// without changing layout.
48    ///
49    /// # See also
50    ///
51    /// [`Tensor::into_owned`] keep data in some conditions, otherwise clone.
52    /// This function can avoid cases where data memory bulk is large, but
53    /// tensor view is small.
54    pub fn into_owned_keep_layout(self) -> Tensor<T, B, D>
55    where
56        R::Data: Clone,
57        R: DataCloneAPI,
58    {
59        let (storage, layout) = self.into_raw_parts();
60        let (data, device) = storage.into_raw_parts();
61        let storage = Storage::new(data.into_owned(), device);
62        unsafe { TensorBase::new_unchecked(storage, layout) }
63    }
64
65    /// Convert tensor into shared tensor.
66    ///
67    /// Data is either moved or cloned.
68    /// Layout is not involved; i.e. all underlying data is moved or cloned
69    /// without changing layout.
70    ///
71    /// # See also
72    ///
73    /// [`Tensor::into_shared`] keep data in some conditions, otherwise clone.
74    /// This function can avoid cases where data memory bulk is large, but
75    /// tensor view is small.
76    pub fn into_shared_keep_layout(self) -> TensorArc<T, B, D>
77    where
78        R::Data: Clone,
79        R: DataCloneAPI,
80    {
81        let (storage, layout) = self.into_raw_parts();
82        let (data, device) = storage.into_raw_parts();
83        let storage = Storage::new(data.into_shared(), device);
84        unsafe { TensorBase::new_unchecked(storage, layout) }
85    }
86}
87
88impl<R, T, B, D> TensorAny<R, T, B, D>
89where
90    R: DataCloneAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
91    R::Data: Clone,
92    D: DimAPI,
93    T: Clone,
94    B: DeviceAPI<T> + DeviceRawAPI<MaybeUninit<T>> + DeviceCreationAnyAPI<T> + OpAssignAPI<T, D>,
95{
96    pub fn into_owned(self) -> Tensor<T, B, D> {
97        let (idx_min, idx_max) = self.layout().bounds_index().rstsr_unwrap();
98        if idx_min == 0 && idx_max == self.storage().len() && idx_max == self.layout().size() {
99            return self.into_owned_keep_layout();
100        } else {
101            return asarray((&self, TensorIterOrder::K));
102        }
103    }
104
105    pub fn into_shared(self) -> TensorArc<T, B, D> {
106        let (idx_min, idx_max) = self.layout().bounds_index().rstsr_unwrap();
107        if idx_min == 0 && idx_max == self.storage().len() && idx_max == self.layout().size() {
108            return self.into_shared_keep_layout();
109        } else {
110            return asarray((&self, TensorIterOrder::K)).into_shared();
111        }
112    }
113
114    pub fn to_owned(&self) -> Tensor<T, B, D> {
115        self.view().into_owned()
116    }
117}
118
119impl<T, B, D> Clone for Tensor<T, B, D>
120where
121    T: Clone,
122    D: DimAPI,
123    B: DeviceAPI<T> + DeviceRawAPI<MaybeUninit<T>> + DeviceCreationAnyAPI<T> + OpAssignAPI<T, D>,
124    <B as DeviceRawAPI<T>>::Raw: Clone,
125{
126    fn clone(&self) -> Self {
127        self.to_owned()
128    }
129}
130
131impl<T, B, D> Clone for TensorCow<'_, T, B, D>
132where
133    T: Clone,
134    D: DimAPI,
135    B: DeviceAPI<T> + DeviceRawAPI<MaybeUninit<T>> + DeviceCreationAnyAPI<T> + OpAssignAPI<T, D>,
136    <B as DeviceRawAPI<T>>::Raw: Clone,
137{
138    fn clone(&self) -> Self {
139        let tsr_owned = self.to_owned();
140        let (storage, layout) = tsr_owned.into_raw_parts();
141        let (data, device) = storage.into_raw_parts();
142        let data = data.into_cow();
143        let storage = Storage::new(data, device);
144        unsafe { TensorBase::new_unchecked(storage, layout) }
145    }
146}
147
148impl<R, T, B, D> TensorAny<R, T, B, D>
149where
150    R: DataAPI<Data = B::Raw> + DataForceMutAPI<B::Raw>,
151    B: DeviceAPI<T>,
152    D: DimAPI,
153{
154    /// # Safety
155    ///
156    /// This function is highly unsafe, as it entirely bypasses Rust's lifetime
157    /// and borrowing rules.
158    pub unsafe fn force_mut(&self) -> TensorMut<'_, T, B, D> {
159        let layout = self.layout().clone();
160        let data = self.data().force_mut();
161        let storage = Storage::new(data, self.device().clone());
162        TensorBase::new_unchecked(storage, layout)
163    }
164}
165
166/* #endregion */
167
168/* #region to_raw */
169
170impl<R, T, B, D> TensorAny<R, T, B, D>
171where
172    R: DataAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
173    T: Clone,
174    D: DimAPI,
175    B: DeviceAPI<T> + DeviceCreationAnyAPI<T> + OpAssignAPI<T, Ix1>,
176{
177    pub fn to_raw_f(&self) -> Result<<B as DeviceRawAPI<T>>::Raw> {
178        rstsr_assert_eq!(self.ndim(), 1, InvalidLayout, "to_vec currently only support 1-D tensor")?;
179        let device = self.device();
180        let layout = self.layout().to_dim::<Ix1>()?;
181        let size = layout.size();
182        let mut new_storage = device.uninit_impl(size)?;
183        device.assign_uninit(new_storage.raw_mut(), &[size].c(), self.raw(), &layout)?;
184        let storage = unsafe { B::assume_init_impl(new_storage) }?;
185        let (data, _) = storage.into_raw_parts();
186        Ok(data.into_raw())
187    }
188
189    pub fn to_vec(&self) -> <B as DeviceRawAPI<T>>::Raw {
190        self.to_raw_f().rstsr_unwrap()
191    }
192}
193
194impl<T, B, D> Tensor<T, B, D>
195where
196    T: Clone,
197    D: DimAPI,
198    B: DeviceAPI<T> + DeviceCreationAnyAPI<T> + OpAssignAPI<T, Ix1>,
199{
200    pub fn into_raw_f(self) -> Result<<B as DeviceRawAPI<T>>::Raw> {
201        rstsr_assert_eq!(self.ndim(), 1, InvalidLayout, "to_vec currently only support 1-D tensor")?;
202        let layout = self.layout();
203        let (idx_min, idx_max) = layout.bounds_index()?;
204        if idx_min == 0 && idx_max == self.storage().len() && idx_max == layout.size() && layout.stride()[0] > 0 {
205            let (storage, _) = self.into_raw_parts();
206            let (data, _) = storage.into_raw_parts();
207            return Ok(data.into_raw());
208        } else {
209            return self.to_raw_f();
210        }
211    }
212
213    pub fn into_raw(self) -> <B as DeviceRawAPI<T>>::Raw {
214        self.into_raw_f().rstsr_unwrap()
215    }
216}
217
218impl<T, B, D> Tensor<T, B, D>
219where
220    T: Clone,
221    D: DimAPI,
222    B: DeviceAPI<T> + DeviceCreationAnyAPI<T> + OpAssignAPI<T, Ix1>,
223    <B as DeviceRawAPI<T>>::Raw: Clone,
224{
225    pub fn into_vec_f(self) -> Result<Vec<T>> {
226        rstsr_assert_eq!(self.ndim(), 1, InvalidLayout, "to_vec currently only support 1-D tensor")?;
227        let layout = self.layout();
228        let (idx_min, idx_max) = layout.bounds_index()?;
229        if idx_min == 0 && idx_max == self.storage().len() && idx_max == layout.size() && layout.stride()[0] > 0 {
230            let (storage, _) = self.into_raw_parts();
231            storage.into_cpu_vec()
232        } else {
233            let data = self.to_raw_f()?;
234            let storage = Storage::new(DataOwned::from(data), self.device().clone());
235            storage.into_cpu_vec()
236        }
237    }
238
239    pub fn into_vec(self) -> Vec<T> {
240        self.into_vec_f().rstsr_unwrap()
241    }
242}
243
244/* #endregion */
245
246/* #region to_scalar */
247
248impl<R, T, B, D> TensorAny<R, T, B, D>
249where
250    R: DataAPI<Data = B::Raw>,
251    T: Clone,
252    D: DimAPI,
253    B: DeviceAPI<T>,
254{
255    pub fn to_scalar_f(&self) -> Result<T> {
256        let layout = self.layout();
257        rstsr_assert_eq!(layout.size(), 1, InvalidLayout)?;
258        // Read the single element at the layout offset directly via `get_index`,
259        // rather than materializing the whole buffer with `to_cpu_vec`. This avoids
260        // a full storage clone for a one-element read and works for non-CPU devices
261        // (which may not expose a cheap CPU `Vec`). Index 0 was previously read,
262        // ignoring slicing/indexing offsets, so e.g. `arange(10).i(9)` (a 0-d view
263        // at offset 9) wrongly returned 0 instead of 9.
264        Ok(self.storage().get_index(layout.offset()))
265    }
266
267    pub fn to_scalar(&self) -> T {
268        self.to_scalar_f().rstsr_unwrap()
269    }
270}
271
272/* #endregion */
273
274/* #region as_ptr */
275
276impl<R, T, B, D> TensorAny<R, T, B, D>
277where
278    R: DataAPI<Data = B::Raw>,
279    D: DimAPI,
280    B: DeviceAPI<T, Raw = Vec<T>>,
281{
282    pub fn as_ptr(&self) -> *const T {
283        unsafe { self.raw().as_ptr().add(self.layout().offset()) }
284    }
285
286    pub fn as_mut_ptr(&mut self) -> *mut T
287    where
288        R: DataMutAPI,
289    {
290        unsafe { self.raw_mut().as_mut_ptr().add(self.layout().offset()) }
291    }
292}
293
294/* #endregion */
295
296/* #region view API */
297
298pub trait TensorViewAPI
299where
300    Self::Dim: DimAPI,
301    Self::Backend: DeviceAPI<Self::Type>,
302{
303    type Type;
304    type Backend;
305    type Dim;
306    /// Get a view of tensor.
307    fn view(&self) -> TensorView<'_, Self::Type, Self::Backend, Self::Dim>;
308}
309
310impl<R, T, B, D> TensorViewAPI for TensorAny<R, T, B, D>
311where
312    D: DimAPI,
313    R: DataAPI<Data = B::Raw>,
314    B: DeviceAPI<T>,
315{
316    type Type = T;
317    type Backend = B;
318    type Dim = D;
319
320    fn view(&self) -> TensorView<'_, T, B, D> {
321        let data = self.data().as_ref();
322        let storage = Storage::new(data, self.device().clone());
323        let layout = self.layout().clone();
324        unsafe { TensorBase::new_unchecked(storage, layout) }
325    }
326}
327
328impl<R, T, B, D> TensorViewAPI for &TensorAny<R, T, B, D>
329where
330    D: DimAPI,
331    R: DataAPI<Data = B::Raw>,
332    B: DeviceAPI<T>,
333{
334    type Type = T;
335    type Backend = B;
336    type Dim = D;
337
338    fn view(&self) -> TensorView<'_, T, B, D> {
339        TensorAny::view(*self)
340    }
341}
342
343impl<R, T, B, D> TensorViewAPI for &mut TensorAny<R, T, B, D>
344where
345    D: DimAPI,
346    R: DataAPI<Data = B::Raw>,
347    B: DeviceAPI<T>,
348{
349    type Type = T;
350    type Backend = B;
351    type Dim = D;
352
353    fn view(&self) -> TensorView<'_, T, B, D> {
354        TensorAny::view(*self)
355    }
356}
357
358pub trait TensorViewMutAPI
359where
360    Self::Dim: DimAPI,
361    Self::Backend: DeviceAPI<Self::Type>,
362{
363    type Type;
364    type Backend;
365    type Dim;
366
367    /// Get a mutable view of tensor.
368    fn view_mut(&mut self) -> TensorMut<'_, Self::Type, Self::Backend, Self::Dim>;
369}
370
371impl<R, T, B, D> TensorViewMutAPI for TensorAny<R, T, B, D>
372where
373    D: DimAPI,
374    R: DataMutAPI<Data = B::Raw>,
375    B: DeviceAPI<T>,
376{
377    type Type = T;
378    type Backend = B;
379    type Dim = D;
380
381    fn view_mut(&mut self) -> TensorMut<'_, T, B, D> {
382        let device = self.device().clone();
383        let layout = self.layout().clone();
384        let data = self.data_mut().as_mut();
385        let storage = Storage::new(data, device);
386        unsafe { TensorBase::new_unchecked(storage, layout) }
387    }
388}
389
390impl<R, T, B, D> TensorViewMutAPI for &mut TensorAny<R, T, B, D>
391where
392    D: DimAPI,
393    R: DataMutAPI<Data = B::Raw>,
394    B: DeviceAPI<T>,
395{
396    type Type = T;
397    type Backend = B;
398    type Dim = D;
399
400    fn view_mut(&mut self) -> TensorMut<'_, T, B, D> {
401        (*self).view_mut()
402    }
403}
404
405pub trait TensorIntoOwnedAPI<T, B, D>
406where
407    D: DimAPI,
408    B: DeviceAPI<T>,
409{
410    /// Convert tensor into owned tensor.
411    ///
412    /// Data is either moved or fully cloned.
413    /// Layout is not involved; i.e. all underlying data is moved or cloned
414    /// without changing layout.
415    fn into_owned(self) -> Tensor<T, B, D>;
416}
417
418impl<R, T, B, D> TensorIntoOwnedAPI<T, B, D> for TensorAny<R, T, B, D>
419where
420    R: DataCloneAPI<Data = <B as DeviceRawAPI<T>>::Raw>,
421    <B as DeviceRawAPI<T>>::Raw: Clone,
422    T: Clone,
423    D: DimAPI,
424    B: DeviceAPI<T> + DeviceRawAPI<MaybeUninit<T>> + DeviceCreationAnyAPI<T> + OpAssignAPI<T, D>,
425{
426    fn into_owned(self) -> Tensor<T, B, D> {
427        TensorAny::into_owned(self)
428    }
429}
430
431/* #endregion */
432
433/* #region tensor prop for computation */
434
435pub trait TensorRefAPI<'l>: TensorViewAPI {}
436impl<'l, R, T, B, D> TensorRefAPI<'l> for &'l TensorAny<R, T, B, D>
437where
438    D: DimAPI,
439    R: DataAPI<Data = B::Raw>,
440    B: DeviceAPI<T>,
441    Self: TensorViewAPI,
442{
443}
444impl<'l, T, B, D> TensorRefAPI<'l> for TensorView<'l, T, B, D>
445where
446    D: DimAPI,
447    B: DeviceAPI<T>,
448    Self: TensorViewAPI,
449{
450}
451
452pub trait TensorRefMutAPI<'l>: TensorViewAPI {}
453impl<'l, R, T, B, D> TensorRefMutAPI<'l> for &mut TensorAny<R, T, B, D>
454where
455    D: DimAPI,
456    R: DataMutAPI<Data = B::Raw>,
457    B: DeviceAPI<T>,
458    Self: TensorViewMutAPI,
459{
460}
461impl<'l, T, B, D> TensorRefMutAPI<'l> for TensorMut<'l, T, B, D>
462where
463    D: DimAPI,
464    B: DeviceAPI<T>,
465    Self: TensorViewMutAPI,
466{
467}
468
469/* #endregion */
470
471#[cfg(test)]
472mod test {
473    use super::*;
474
475    #[test]
476    fn test_into_cow() {
477        let mut a = arange(3);
478        let ptr_a = a.raw().as_ptr();
479
480        let a_mut = a.view_mut();
481        let a_cow = a_mut.into_cow();
482        println!("{a_cow:?}");
483
484        let a_ref = a.view();
485        let a_cow = a_ref.into_cow();
486        println!("{a_cow:?}");
487
488        let a_cow = a.into_cow();
489        println!("{a_cow:?}");
490        let ptr_a_cow = a_cow.raw().as_ptr();
491        assert_eq!(ptr_a, ptr_a_cow);
492    }
493
494    #[test]
495    #[ignore]
496    fn test_force_mut() {
497        let n = 4096;
498        let a = linspace((0.0, 1.0, n * n)).into_shape((n, n));
499        for _ in 0..10 {
500            let time = std::time::Instant::now();
501            for i in 0..n {
502                let a_view = a.slice(i);
503                let mut a_mut = unsafe { a_view.force_mut() };
504                a_mut *= i as f64 / 2048.0;
505            }
506            println!("Elapsed time {:?}", time.elapsed());
507        }
508        println!("{a:16.10}");
509    }
510
511    #[test]
512    #[ignore]
513    #[cfg(feature = "rayon")]
514    fn test_force_mut_par() {
515        use rayon::prelude::*;
516        let n = 4096;
517        let a = linspace((0.0, 1.0, n * n)).into_shape((n, n));
518        for _ in 0..10 {
519            let time = std::time::Instant::now();
520            (0..n).into_par_iter().for_each(|i| {
521                let a_view = a.slice(i);
522                let mut a_mut = unsafe { a_view.force_mut() };
523                a_mut *= i as f64 / 2048.0;
524            });
525            println!("Elapsed time {:?}", time.elapsed());
526        }
527        println!("{a:16.10}");
528    }
529}
rstsr_core/tensor/ownership_conversion.rs

rstsr_core/tensor/
ownership_conversion.rs
