xtrain/crates/xtrain-tensor/src/storage.rs

//! Tensor storage: host (CPU) bytes or a GPU buffer, reference-counted so
//! views (clones with different shape/strides) can share the backing data.

use std::sync::Arc;
use xtrain_cuda::{GpuBuffer, Result as CudaResult};

enum StorageInner {
    Cpu { data: Vec<u8> },
    Cuda { buffer: GpuBuffer, device: u32 },
}

/// Reference-counted tensor storage. Cloning is cheap (bumps the `Arc`).
#[derive(Clone)]
pub struct Storage(Arc<StorageInner>);

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Device {
    Cpu,
    Cuda(u32),
}

impl std::fmt::Display for Device {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Device::Cpu => write!(f, "cpu"),
            Device::Cuda(i) => write!(f, "cuda:{i}"),
        }
    }
}

impl Storage {
    pub fn cpu(data: Vec<u8>) -> Self {
        Self(Arc::new(StorageInner::Cpu { data }))
    }

    pub fn cuda(buffer: GpuBuffer, device: u32) -> Self {
        Self(Arc::new(StorageInner::Cuda { buffer, device }))
    }

    pub fn device(&self) -> Device {
        match self.0.as_ref() {
            StorageInner::Cpu { .. } => Device::Cpu,
            StorageInner::Cuda { device, .. } => Device::Cuda(*device),
        }
    }

    pub fn len_bytes(&self) -> usize {
        match self.0.as_ref() {
            StorageInner::Cpu { data } => data.len(),
            StorageInner::Cuda { buffer, .. } => buffer.len(),
        }
    }

    /// Read-only view of CPU bytes. Panics if the storage lives on the GPU.
    pub fn as_cpu_bytes(&self) -> &[u8] {
        match self.0.as_ref() {
            StorageInner::Cpu { data } => data,
            StorageInner::Cuda { .. } => panic!("cannot read GPU storage as CPU bytes"),
        }
    }

    /// Borrow the GPU buffer. Panics if the storage lives on the CPU.
    pub fn gpu_buffer(&self) -> &GpuBuffer {
        match self.0.as_ref() {
            StorageInner::Cuda { buffer, .. } => buffer,
            StorageInner::Cpu { .. } => panic!("cannot read CPU storage as GPU buffer"),
        }
    }

    /// Copy to another device. Returns a clone of the `Arc` when already there.
    /// T2 supports CPU↔CUDA(0); device-to-device copy across GPUs is out of scope.
    pub fn to_device(&self, target: Device) -> CudaResult<Self> {
        let current = self.device();
        if current == target {
            return Ok(self.clone());
        }
        match (current, target) {
            (Device::Cpu, Device::Cuda(dev)) => {
                let host = self.as_cpu_bytes();
                let mut buf = GpuBuffer::alloc(host.len())?;
                buf.copy_from_host(host)?;
                Ok(Storage::cuda(buf, dev))
            }
            (Device::Cuda(_), Device::Cpu) => {
                let src = self.gpu_buffer();
                let mut host = vec![0u8; src.len()];
                src.copy_to_host(&mut host)?;
                Ok(Storage::cpu(host))
            }
            (Device::Cuda(_), Device::Cuda(_)) => {
                panic!("cross-GPU storage transfer is not supported in T2")
            }
            _ => unreachable!(),
        }
    }

    /// Zeroed storage on the given device.
    pub fn zeros(len_bytes: usize, device: Device) -> CudaResult<Self> {
        match device {
            Device::Cpu => Ok(Storage::cpu(vec![0u8; len_bytes])),
            Device::Cuda(dev) => {
                // Device-side memset (Phase T7): avoids a blocking H2D memcpy of a
                // host zero buffer on every op-output allocation. cudaMemset is
                // async on the default stream, so it doesn't serialize the stream.
                let mut buf = GpuBuffer::alloc(len_bytes)?;
                buf.memset(0)?;
                Ok(Storage::cuda(buf, dev))
            }
        }
    }
}