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

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

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

/// Reference-counted storage for tensor data. Multiple tensors can share
/// the same storage (e.g., after transpose or slice — view semantics).
#[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(),
        }
    }

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

    pub fn gpu_buffer(&self) -> &GpuBuffer {
        match self.0.as_ref() {
            StorageInner::Cuda { buffer, .. } => buffer,
            StorageInner::Cpu { .. } => panic!("cannot access CPU storage as GPU buffer"),
        }
    }

    /// Copy to a different device. If already on the target device, clones the Arc (no copy).
    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 cpu_data = self.as_cpu_bytes();
                let mut buf = xserv_cuda::allocator::cached_alloc(cpu_data.len())?;
                buf.copy_from_host(cpu_data)?;
                Ok(Storage::cuda(buf, dev))
            }
            (Device::Cuda(_), Device::Cpu) => {
                let gpu_buf = self.gpu_buffer();
                let mut data = vec![0u8; gpu_buf.len()];
                gpu_buf.copy_to_host(&mut data)?;
                Ok(Storage::cpu(data))
            }
            (Device::Cuda(_), Device::Cuda(dev)) => {
                let src = self.gpu_buffer();
                let mut dst = xserv_cuda::allocator::cached_alloc(src.len())?;
                dst.copy_from_device(src)?;
                Ok(Storage::cuda(dst, dev))
            }
            _ => unreachable!(),
        }
    }

    /// Create a new owned copy of the storage on the same device.
    pub fn deep_copy(&self) -> CudaResult<Self> {
        match self.0.as_ref() {
            StorageInner::Cpu { data } => Ok(Storage::cpu(data.clone())),
            StorageInner::Cuda { buffer, device } => {
                let mut dst = xserv_cuda::allocator::cached_alloc(buffer.len())?;
                dst.copy_from_device(buffer)?;
                Ok(Storage::cuda(dst, *device))
            }
        }
    }

    /// Allocate 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) => {
                let mut buf = xserv_cuda::allocator::cached_alloc(len_bytes)?;
                buf.zero()?;
                Ok(Storage::cuda(buf, dev))
            }
        }
    }

    /// Allocate storage **without zeroing** on the given device.
    /// The buffer may contain stale data from the caching allocator's pool.
    /// Only use when the caller guarantees the kernel will fully overwrite
    /// every element before any read.
    pub fn empty(len_bytes: usize, device: Device) -> CudaResult<Self> {
        match device {
            Device::Cpu => Ok(Storage::cpu(vec![0u8; len_bytes])), // CPU still zeros (cheap)
            Device::Cuda(dev) => {
                let buf = xserv_cuda::allocator::cached_alloc(len_bytes)?;
                Ok(Storage::cuda(buf, dev))
            }
        }
    }
}