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,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]);
}
}