diff --git a/crates/yscv-tensor/src/shape.rs b/crates/yscv-tensor/src/shape.rs index 044dc78..6f5f544 100644 --- a/crates/yscv-tensor/src/shape.rs +++ b/crates/yscv-tensor/src/shape.rs @@ -1,17 +1,32 @@ +use super::tensor::{DimsVec, INLINE_CAP}; + pub(crate) fn shape_element_count(shape: &[usize]) -> Option { shape .iter() .try_fold(1usize, |acc, dim| acc.checked_mul(*dim)) } -pub(crate) fn compute_strides(shape: &[usize]) -> Option> { +pub(crate) fn compute_strides(shape: &[usize]) -> Option { + if shape.len() <= INLINE_CAP { + let mut buf = [0usize; INLINE_CAP]; + let mut stride = 1usize; + for axis in (0..shape.len()).rev() { + buf[axis] = stride; + stride = stride.checked_mul(shape[axis])?; + } + return Some(DimsVec::Inline { + buf, + len: shape.len() as u8, + }); + } + let mut strides = vec![0usize; shape.len()]; let mut stride = 1usize; for axis in (0..shape.len()).rev() { strides[axis] = stride; stride = stride.checked_mul(shape[axis])?; } - Some(strides) + Some(DimsVec::Heap(strides)) } pub(crate) fn broadcast_shape(left: &[usize], right: &[usize]) -> Option> { diff --git a/crates/yscv-tensor/src/tensor.rs b/crates/yscv-tensor/src/tensor.rs index 6b37a06..7b07e22 100644 --- a/crates/yscv-tensor/src/tensor.rs +++ b/crates/yscv-tensor/src/tensor.rs @@ -7,7 +7,7 @@ use super::shape::{compute_strides, shape_element_count}; // ── Inline shape/strides storage (no heap alloc for ≤6D tensors) ───────── // WHY 6: covers all common tensor ranks (scalar(0)..conv weight(5)) without heap allocation. -const INLINE_CAP: usize = 6; +pub(crate) const INLINE_CAP: usize = 6; /// Stack-allocated small vector for tensor shape/strides. /// Stores up to 6 dimensions inline; falls back to heap for higher ranks. @@ -267,7 +267,7 @@ impl Tensor { Ok(Self { shape: DimsVec::from(shape), - strides: DimsVec::from(strides), + strides, storage: Arc::new(Storage::F32(data)), device: Device::Cpu, layout: Layout::NCHW, @@ -292,7 +292,7 @@ impl Tensor { Ok(Self { shape: DimsVec::from(shape), - strides: DimsVec::from(strides), + strides, storage: Arc::new(Storage::F32(AlignedVec::from_vec(data))), device: Device::Cpu, layout: Layout::NCHW, @@ -315,7 +315,7 @@ impl Tensor { })?; Ok(Self { shape: DimsVec::from(shape), - strides: DimsVec::from(strides), + strides, storage: Arc::new(Storage::F16(data)), device: Device::Cpu, layout: Layout::NCHW, @@ -338,7 +338,7 @@ impl Tensor { })?; Ok(Self { shape: DimsVec::from(shape), - strides: DimsVec::from(strides), + strides, storage: Arc::new(Storage::BF16(data)), device: Device::Cpu, layout: Layout::NCHW, @@ -370,7 +370,7 @@ impl Tensor { Ok(Self { shape: DimsVec::from(shape), - strides: DimsVec::from(strides), + strides, storage: Arc::new(Storage::F32(AlignedVec::filled(count, value))), device: Device::Cpu, layout: Layout::NCHW, @@ -391,7 +391,7 @@ impl Tensor { Ok(Self { shape: DimsVec::from(shape), - strides: DimsVec::from(strides), + strides, storage: Arc::new(Storage::F32(AlignedVec::calloc(count))), device: Device::Cpu, layout: Layout::NCHW, @@ -631,7 +631,7 @@ impl Tensor { Ok(Self { shape: DimsVec::from(new_shape), - strides: DimsVec::from(new_strides), + strides: new_strides, storage: self.storage.clone(), device: self.device, layout: self.layout, @@ -657,7 +657,7 @@ impl Tensor { Ok(Self { shape: DimsVec::from(new_shape), - strides: DimsVec::from(new_strides), + strides: new_strides, storage: self.storage, device: self.device, layout: self.layout,