luma_tensor/ops/
transfer.rs

1//! Cross-device transfer: move a tensor to another device, e.g. `Cpu` → `Cuda`.
2
3use std::any::TypeId;
4
5#[cfg(feature = "cuda")]
6use crate::Cuda;
7use crate::ops::construct::BytesDTypeKind;
8use crate::{Bool, Cpu, DTypeKind, Device, Float, Int, Tensor};
9
10// ============================================================================
11//    TransferDTypeKind: kind dispatch for `to_device`
12// ============================================================================
13
14/// Kind-level dispatch for `Tensor::to_device`.
15///
16/// Implemented for the three closed kinds ([`Float`], [`Int`], [`Bool`]) for
17/// every source/target device pair. The transfer is a bytes round-trip in
18/// logical order (`to_bytes` → `from_bytes`), so the result is contiguous.
19pub trait TransferDTypeKind<D: Device, D2: Device>: DTypeKind<D> + DTypeKind<D2> + BytesDTypeKind<D> + BytesDTypeKind<D2> {
20    fn transfer(src: &Tensor<D, Self>, device: &D2) -> crate::Result<Tensor<D2, Self>>;
21}
22
23impl<D: Device, D2: Device> TransferDTypeKind<D, D2> for Float {
24    fn transfer(src: &Tensor<D, Float>, device: &D2) -> crate::Result<Tensor<D2, Float>> {
25        let bytes = src.to_bytes()?;
26        let out = Tensor::from_bytes(bytes, src.shape().clone(), (device, src.dtype()))?;
27        // The autograd graph is single-device (`Op<D>` holds same-device
28        // inputs), so the transfer cannot record a node: the result is a fresh
29        // leaf. Preserve the trainability flag so a moved parameter keeps
30        // accumulating gradients.
31        out.set_requires_grad(src.requires_grad());
32        Ok(out)
33    }
34}
35
36impl<D: Device, D2: Device> TransferDTypeKind<D, D2> for Int {
37    fn transfer(src: &Tensor<D, Int>, device: &D2) -> crate::Result<Tensor<D2, Int>> {
38        let bytes = src.to_bytes()?;
39        Tensor::from_bytes(bytes, src.shape().clone(), (device, src.dtype()))
40    }
41}
42
43impl<D: Device, D2: Device> TransferDTypeKind<D, D2> for Bool {
44    fn transfer(src: &Tensor<D, Bool>, device: &D2) -> crate::Result<Tensor<D2, Bool>> {
45        let bytes = src.to_bytes()?;
46        Tensor::from_bytes(bytes, src.shape().clone(), (device, src.dtype()))
47    }
48}
49
50// ============================================================================
51//    Public API
52// ============================================================================
53
54impl<D: Device, K: DTypeKind<D>> Tensor<D, K> {
55    /// Move the tensor to another device.
56    ///
57    /// - Same device: returns this handle unchanged (shared `Arc`, O(1)) —
58    ///   generic code can call this unconditionally.
59    /// - Cross device: deep-copies the data (`to_bytes`/`from_bytes`), dtype
60    ///   unchanged, result contiguous. For `Float` tensors the autograd graph
61    ///   is severed (the result is a leaf) but `requires_grad` is preserved.
62    /// - Meta tensors (no storage) error with [`Error::MetaTensor`](crate::Error::MetaTensor).
63    pub fn to_device<D2: Device>(&self, device: &D2) -> crate::Result<Tensor<D2, K>>
64    where
65        K: TransferDTypeKind<D, D2>,
66    {
67        if TypeId::of::<D>() == TypeId::of::<D2>() {
68            // SAFETY: `TypeId` equality means `D` and `D2` are the same concrete
69            // type (`Device: 'static`), so these casts are type puns between
70            // identical types with identical layout — no aliasing or validity
71            // concerns beyond a plain clone.
72            let target: &D = unsafe { &*(device as *const D2 as *const D) };
73            if self.device().same_device(target) {
74                let p = self as *const Tensor<D, K> as *const Tensor<D2, K>;
75                return Ok(unsafe { (*p).clone() });
76            }
77        }
78        // Cross-device copy (covers `Cuda` → `Cuda` with different ordinals,
79        // which goes through the host — there is no DtoD primitive yet).
80        K::transfer(self, device)
81    }
82
83    /// Convenience for `to_device(&Cpu)`.
84    pub fn cpu(&self) -> crate::Result<Tensor<Cpu, K>>
85    where
86        K: TransferDTypeKind<D, Cpu>,
87    {
88        self.to_device(&Cpu)
89    }
90
91    /// Convenience for `to_device(&Cuda::new(ordinal))`.
92    #[cfg(feature = "cuda")]
93    pub fn cuda(&self, ordinal: usize) -> crate::Result<Tensor<Cuda, K>>
94    where
95        K: TransferDTypeKind<D, Cuda>,
96    {
97        self.to_device(&Cuda::new(ordinal)?)
98    }
99}
luma_tensor/ops/transfer.rs

luma_tensor/ops/
transfer.rs
