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:
119
crates/xserv-tensor/src/storage.rs
Normal file
119
crates/xserv-tensor/src/storage.rs
Normal 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))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user