phase 2: tensor abstraction layer

- DType enum (F32, F16, BF16) with TensorDType trait
- Shape utilities: contiguous_strides, broadcast_shape, broadcast_strides
- Storage with Arc reference counting (CPU Vec<u8> or GPU GpuBuffer)
- Device enum (Cpu, Cuda(id)) with to_device transfer
- Tensor type with strided layout: reshape, transpose, squeeze, unsqueeze
- contiguous() copies non-contiguous views to contiguous layout
- from_slice, zeros, ones constructors
- as_slice<T> for typed CPU read access, data_ptr for GPU kernel launch
- CPU↔GPU roundtrip verified
- All 27 tests pass (12 cuda + 4 shape + 11 tensor)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-05-21 19:45:22 +08:00
parent c8f7bc0c3c
commit a83971fa25
8 changed files with 654 additions and 0 deletions

View File

@@ -0,0 +1,119 @@
use std::sync::Arc;
use xserv_cuda::{GpuBuffer, Result as CudaResult};
enum StorageInner {
Cpu { data: Vec<u8> },
Cuda { buffer: GpuBuffer },
}
/// 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) -> Self {
Self(Arc::new(StorageInner::Cuda { buffer }))
}
pub fn device(&self) -> Device {
match self.0.as_ref() {
StorageInner::Cpu { .. } => Device::Cpu,
StorageInner::Cuda { .. } => Device::Cuda(0),
}
}
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 = GpuBuffer::alloc(cpu_data.len())?;
buf.copy_from_host(cpu_data)?;
Ok(Storage::cuda(buf))
}
(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(_)) => {
let src = self.gpu_buffer();
let mut dst = GpuBuffer::alloc(src.len())?;
dst.copy_from_device(src)?;
Ok(Storage::cuda(dst))
}
_ => 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 } => {
let mut dst = GpuBuffer::alloc(buffer.len())?;
dst.copy_from_device(buffer)?;
Ok(Storage::cuda(dst))
}
}
}
/// 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(_) => {
let mut buf = GpuBuffer::alloc(len_bytes)?;
buf.zero()?;
Ok(Storage::cuda(buf))
}
}
}
}