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:
105
crates/xserv-tensor/src/shape.rs
Normal file
105
crates/xserv-tensor/src/shape.rs
Normal file
@@ -0,0 +1,105 @@
|
||||
use smallvec::SmallVec;
|
||||
|
||||
pub type Dims = SmallVec<[usize; 4]>;
|
||||
|
||||
/// Compute contiguous strides for a given shape (row-major / C order).
|
||||
/// Example: shape [2, 3, 4] => strides [12, 4, 1]
|
||||
pub fn contiguous_strides(shape: &[usize]) -> Dims {
|
||||
let mut strides = SmallVec::with_capacity(shape.len());
|
||||
strides.resize(shape.len(), 0);
|
||||
if shape.is_empty() {
|
||||
return strides;
|
||||
}
|
||||
strides[shape.len() - 1] = 1;
|
||||
for i in (0..shape.len() - 1).rev() {
|
||||
strides[i] = strides[i + 1] * shape[i + 1];
|
||||
}
|
||||
strides
|
||||
}
|
||||
|
||||
/// Check if the given strides represent contiguous (row-major) layout for the shape.
|
||||
pub fn is_contiguous(shape: &[usize], strides: &[usize]) -> bool {
|
||||
if shape.is_empty() {
|
||||
return true;
|
||||
}
|
||||
let expected = contiguous_strides(shape);
|
||||
strides == expected.as_slice()
|
||||
}
|
||||
|
||||
/// Total number of elements given a shape.
|
||||
pub fn num_elements(shape: &[usize]) -> usize {
|
||||
shape.iter().product()
|
||||
}
|
||||
|
||||
/// Compute the shape after broadcasting two shapes together (NumPy rules).
|
||||
/// Returns None if shapes are not broadcastable.
|
||||
pub fn broadcast_shape(a: &[usize], b: &[usize]) -> Option<Dims> {
|
||||
let ndim = a.len().max(b.len());
|
||||
let mut result = SmallVec::with_capacity(ndim);
|
||||
for i in 0..ndim {
|
||||
let da = if i < ndim - a.len() { 1 } else { a[i - (ndim - a.len())] };
|
||||
let db = if i < ndim - b.len() { 1 } else { b[i - (ndim - b.len())] };
|
||||
if da == db {
|
||||
result.push(da);
|
||||
} else if da == 1 {
|
||||
result.push(db);
|
||||
} else if db == 1 {
|
||||
result.push(da);
|
||||
} else {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
Some(result)
|
||||
}
|
||||
|
||||
/// Compute broadcast strides: for dimensions where size is 1 but output is >1, stride becomes 0.
|
||||
pub fn broadcast_strides(shape: &[usize], strides: &[usize], target_shape: &[usize]) -> Dims {
|
||||
let ndim = target_shape.len();
|
||||
let offset = ndim - shape.len();
|
||||
let mut result = SmallVec::with_capacity(ndim);
|
||||
for i in 0..ndim {
|
||||
if i < offset {
|
||||
result.push(0);
|
||||
} else {
|
||||
let orig_idx = i - offset;
|
||||
if shape[orig_idx] == 1 && target_shape[i] > 1 {
|
||||
result.push(0);
|
||||
} else {
|
||||
result.push(strides[orig_idx]);
|
||||
}
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_contiguous_strides() {
|
||||
assert_eq!(contiguous_strides(&[2, 3, 4]).as_slice(), &[12, 4, 1]);
|
||||
assert_eq!(contiguous_strides(&[5]).as_slice(), &[1]);
|
||||
assert_eq!(contiguous_strides(&[2, 3]).as_slice(), &[3, 1]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_contiguous() {
|
||||
assert!(is_contiguous(&[2, 3], &[3, 1]));
|
||||
assert!(!is_contiguous(&[3, 2], &[1, 3])); // transposed
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_broadcast_shape() {
|
||||
assert_eq!(broadcast_shape(&[3, 1], &[1, 4]).unwrap().as_slice(), &[3, 4]);
|
||||
assert_eq!(broadcast_shape(&[2, 3, 4], &[4]).unwrap().as_slice(), &[2, 3, 4]);
|
||||
assert_eq!(broadcast_shape(&[1], &[5, 3]).unwrap().as_slice(), &[5, 3]);
|
||||
assert!(broadcast_shape(&[3], &[4]).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_broadcast_strides() {
|
||||
// [3,1] with strides [1,1] broadcast to [3,4]
|
||||
assert_eq!(broadcast_strides(&[3, 1], &[1, 1], &[3, 4]).as_slice(), &[1, 0]);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user