1ab6ca9c09
narrow(dim, start, len) creates a zero-copy slice along any dimension. is_contiguous() now ignores stride mismatches on dimensions of size 1, since those dimensions are never stepped. This avoids unnecessary GPU strided copies when slicing fused projection outputs at batch=1. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
115 lines
3.5 KiB
Rust
115 lines
3.5 KiB
Rust
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.
|
|
/// A stride mismatch on a dimension of size 1 is allowed because that
|
|
/// dimension is never stepped.
|
|
pub fn is_contiguous(shape: &[usize], strides: &[usize]) -> bool {
|
|
if shape.is_empty() {
|
|
return true;
|
|
}
|
|
let ndim = shape.len();
|
|
let mut expected_stride = 1usize;
|
|
for d in (0..ndim).rev() {
|
|
if shape[d] != 1 && strides[d] != expected_stride {
|
|
return false;
|
|
}
|
|
expected_stride *= shape[d];
|
|
}
|
|
true
|
|
}
|
|
|
|
/// 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]);
|
|
}
|
|
}
|