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:
228
crates/xserv-tensor/src/tensor.rs
Normal file
228
crates/xserv-tensor/src/tensor.rs
Normal file
@@ -0,0 +1,228 @@
|
||||
use crate::dtype::{DType, TensorDType};
|
||||
use crate::shape::{self, Dims};
|
||||
use crate::storage::{Device, Storage};
|
||||
|
||||
/// Multi-dimensional array with CPU or GPU storage.
|
||||
///
|
||||
/// Tensors support view semantics: transpose, slice, etc. share
|
||||
/// the underlying storage and only change shape/strides/offset.
|
||||
#[derive(Clone)]
|
||||
pub struct Tensor {
|
||||
storage: Storage,
|
||||
shape: Dims,
|
||||
strides: Dims,
|
||||
offset: usize,
|
||||
dtype: DType,
|
||||
}
|
||||
|
||||
impl Tensor {
|
||||
// --- Creation ---
|
||||
|
||||
pub fn from_slice<T: TensorDType>(data: &[T], shape: &[usize]) -> Self {
|
||||
let numel: usize = shape.iter().product();
|
||||
assert_eq!(data.len(), numel, "data length mismatch with shape");
|
||||
let bytes = unsafe {
|
||||
std::slice::from_raw_parts(data.as_ptr() as *const u8, numel * T::DTYPE.size_bytes())
|
||||
};
|
||||
Self {
|
||||
storage: Storage::cpu(bytes.to_vec()),
|
||||
shape: Dims::from_slice(shape),
|
||||
strides: shape::contiguous_strides(shape),
|
||||
offset: 0,
|
||||
dtype: T::DTYPE,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn zeros(shape: &[usize], dtype: DType, device: Device) -> Self {
|
||||
let numel = shape::num_elements(shape);
|
||||
let len_bytes = numel * dtype.size_bytes();
|
||||
let storage = Storage::zeros(len_bytes, device).expect("alloc failed");
|
||||
Self {
|
||||
storage,
|
||||
shape: Dims::from_slice(shape),
|
||||
strides: shape::contiguous_strides(shape),
|
||||
offset: 0,
|
||||
dtype,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn ones(shape: &[usize], dtype: DType) -> Self {
|
||||
let numel = shape::num_elements(shape);
|
||||
match dtype {
|
||||
DType::F32 => Self::from_slice(&vec![1.0f32; numel], shape),
|
||||
DType::F16 => Self::from_slice(&vec![half::f16::from_f32(1.0); numel], shape),
|
||||
DType::BF16 => Self::from_slice(&vec![half::bf16::from_f32(1.0); numel], shape),
|
||||
}
|
||||
}
|
||||
|
||||
// --- Properties ---
|
||||
|
||||
pub fn shape(&self) -> &[usize] { &self.shape }
|
||||
pub fn strides(&self) -> &[usize] { &self.strides }
|
||||
pub fn dtype(&self) -> DType { self.dtype }
|
||||
pub fn ndim(&self) -> usize { self.shape.len() }
|
||||
pub fn numel(&self) -> usize { shape::num_elements(&self.shape) }
|
||||
pub fn offset(&self) -> usize { self.offset }
|
||||
|
||||
pub fn device(&self) -> Device { self.storage.device() }
|
||||
|
||||
pub fn is_contiguous(&self) -> bool {
|
||||
shape::is_contiguous(&self.shape, &self.strides)
|
||||
}
|
||||
|
||||
// --- Shape operations (view, no copy) ---
|
||||
|
||||
pub fn reshape(&self, new_shape: &[usize]) -> Self {
|
||||
assert!(self.is_contiguous(), "reshape requires contiguous tensor");
|
||||
let new_numel: usize = new_shape.iter().product();
|
||||
assert_eq!(new_numel, self.numel(), "reshape numel mismatch");
|
||||
Self {
|
||||
storage: self.storage.clone(),
|
||||
shape: Dims::from_slice(new_shape),
|
||||
strides: shape::contiguous_strides(new_shape),
|
||||
offset: self.offset,
|
||||
dtype: self.dtype,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn transpose(&self, dim0: usize, dim1: usize) -> Self {
|
||||
assert!(dim0 < self.ndim() && dim1 < self.ndim());
|
||||
let mut new_shape = self.shape.clone();
|
||||
let mut new_strides = self.strides.clone();
|
||||
new_shape.swap(dim0, dim1);
|
||||
new_strides.swap(dim0, dim1);
|
||||
Self {
|
||||
storage: self.storage.clone(),
|
||||
shape: new_shape,
|
||||
strides: new_strides,
|
||||
offset: self.offset,
|
||||
dtype: self.dtype,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn squeeze(&self, dim: usize) -> Self {
|
||||
assert!(dim < self.ndim() && self.shape[dim] == 1);
|
||||
let mut new_shape = self.shape.clone();
|
||||
let mut new_strides = self.strides.clone();
|
||||
new_shape.remove(dim);
|
||||
new_strides.remove(dim);
|
||||
Self {
|
||||
storage: self.storage.clone(),
|
||||
shape: new_shape,
|
||||
strides: new_strides,
|
||||
offset: self.offset,
|
||||
dtype: self.dtype,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn unsqueeze(&self, dim: usize) -> Self {
|
||||
assert!(dim <= self.ndim());
|
||||
let mut new_shape = self.shape.clone();
|
||||
let mut new_strides = self.strides.clone();
|
||||
new_shape.insert(dim, 1);
|
||||
let stride_val = if dim < self.strides.len() { self.strides[dim] } else { 1 };
|
||||
new_strides.insert(dim, stride_val);
|
||||
Self {
|
||||
storage: self.storage.clone(),
|
||||
shape: new_shape,
|
||||
strides: new_strides,
|
||||
offset: self.offset,
|
||||
dtype: self.dtype,
|
||||
}
|
||||
}
|
||||
|
||||
/// Make contiguous: if already contiguous, return clone (shared storage).
|
||||
/// Otherwise, copy data into a new contiguous buffer.
|
||||
pub fn contiguous(&self) -> Self {
|
||||
if self.is_contiguous() {
|
||||
return self.clone();
|
||||
}
|
||||
// Copy to contiguous layout on CPU
|
||||
assert_eq!(self.device(), Device::Cpu, "contiguous() on GPU not yet supported");
|
||||
let numel = self.numel();
|
||||
let elem_size = self.dtype.size_bytes();
|
||||
let src_bytes = self.storage.as_cpu_bytes();
|
||||
let mut dst = vec![0u8; numel * elem_size];
|
||||
// Iterate all elements using strides
|
||||
let ndim = self.ndim();
|
||||
let mut idx = vec![0usize; ndim];
|
||||
for flat in 0..numel {
|
||||
let src_offset = self.offset + idx.iter().zip(self.strides.iter()).map(|(i, s)| i * s).sum::<usize>();
|
||||
let src_byte_offset = src_offset * elem_size;
|
||||
let dst_byte_offset = flat * elem_size;
|
||||
dst[dst_byte_offset..dst_byte_offset + elem_size]
|
||||
.copy_from_slice(&src_bytes[src_byte_offset..src_byte_offset + elem_size]);
|
||||
// Increment index (rightmost first)
|
||||
for d in (0..ndim).rev() {
|
||||
idx[d] += 1;
|
||||
if idx[d] < self.shape[d] {
|
||||
break;
|
||||
}
|
||||
idx[d] = 0;
|
||||
}
|
||||
}
|
||||
Self {
|
||||
storage: Storage::cpu(dst),
|
||||
shape: self.shape.clone(),
|
||||
strides: shape::contiguous_strides(&self.shape),
|
||||
offset: 0,
|
||||
dtype: self.dtype,
|
||||
}
|
||||
}
|
||||
|
||||
// --- Device transfer ---
|
||||
|
||||
pub fn to_device(&self, device: Device) -> Self {
|
||||
let t = if self.is_contiguous() { self.clone() } else { self.contiguous() };
|
||||
if t.device() == device {
|
||||
return t;
|
||||
}
|
||||
let new_storage = t.storage.to_device(device).expect("device transfer failed");
|
||||
Self {
|
||||
storage: new_storage,
|
||||
shape: t.shape,
|
||||
strides: t.strides,
|
||||
offset: 0,
|
||||
dtype: t.dtype,
|
||||
}
|
||||
}
|
||||
|
||||
// --- Data access (CPU only) ---
|
||||
|
||||
/// Read tensor data as a typed slice. Requires contiguous CPU tensor.
|
||||
pub fn as_slice<T: TensorDType>(&self) -> &[T] {
|
||||
assert_eq!(T::DTYPE, self.dtype, "dtype mismatch");
|
||||
assert!(self.is_contiguous(), "as_slice requires contiguous");
|
||||
assert_eq!(self.device(), Device::Cpu, "as_slice requires CPU");
|
||||
let bytes = self.storage.as_cpu_bytes();
|
||||
let elem_size = self.dtype.size_bytes();
|
||||
let start = self.offset * elem_size;
|
||||
let len = self.numel();
|
||||
unsafe { std::slice::from_raw_parts(bytes[start..].as_ptr() as *const T, len) }
|
||||
}
|
||||
|
||||
/// Raw pointer to storage start (for GPU kernel launch).
|
||||
pub fn data_ptr(&self) -> *const u8 {
|
||||
match self.device() {
|
||||
Device::Cpu => {
|
||||
let bytes = self.storage.as_cpu_bytes();
|
||||
unsafe { bytes.as_ptr().add(self.offset * self.dtype.size_bytes()) }
|
||||
}
|
||||
Device::Cuda(_) => {
|
||||
let buf = self.storage.gpu_buffer();
|
||||
unsafe { buf.as_ptr().add(self.offset * self.dtype.size_bytes()) }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn storage(&self) -> &Storage { &self.storage }
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for Tensor {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(
|
||||
f, "Tensor(shape={:?}, dtype={}, device={}, contiguous={})",
|
||||
self.shape.as_slice(), self.dtype, self.device(), self.is_contiguous()
|
||||
)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user