Gahow Wang
fbd07a578c
New xtrain-tensor crate: DType (F32), shape/stride helpers, Arc-counted host/device Storage with CPU↔CUDA copy, and a contiguous Tensor with creation, host↔device transfer, and a scale() op driving the elementwise kernel. GPU integration tests (host↔device roundtrip + scale correctness) gated behind not(no_cuda); a thin build.rs emits the no_cuda cfg so the kernel call sites compile out locally. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
58 lines
1.6 KiB
Rust
58 lines
1.6 KiB
Rust
//! Shape / stride helpers. Strides are in **elements** (not bytes), row-major.
|
|
|
|
use smallvec::SmallVec;
|
|
|
|
/// Inline storage for the common ≤4D case; spills to the heap beyond that.
|
|
pub type Dims = SmallVec<[usize; 4]>;
|
|
|
|
/// Row-major (C order) contiguous strides for a shape.
|
|
/// Example: `[2, 3, 4]` => `[12, 4, 1]`.
|
|
pub fn contiguous_strides(shape: &[usize]) -> Dims {
|
|
let mut strides: Dims = 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
|
|
}
|
|
|
|
/// True if `strides` describe a row-major contiguous layout for `shape`.
|
|
/// A mismatched stride on a size-1 dimension is fine (it is never stepped).
|
|
pub fn is_contiguous(shape: &[usize], strides: &[usize]) -> bool {
|
|
let ndim = shape.len();
|
|
let mut expected = 1usize;
|
|
for d in (0..ndim).rev() {
|
|
if shape[d] != 1 && strides[d] != expected {
|
|
return false;
|
|
}
|
|
expected *= shape[d];
|
|
}
|
|
true
|
|
}
|
|
|
|
/// Total element count.
|
|
pub fn num_elements(shape: &[usize]) -> usize {
|
|
shape.iter().product()
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn contiguous_strides_basic() {
|
|
assert_eq!(contiguous_strides(&[2, 3, 4]).as_slice(), &[12, 4, 1]);
|
|
assert_eq!(contiguous_strides(&[5]).as_slice(), &[1]);
|
|
}
|
|
|
|
#[test]
|
|
fn is_contiguous_detects_transpose() {
|
|
assert!(is_contiguous(&[2, 3], &[3, 1]));
|
|
assert!(!is_contiguous(&[3, 2], &[1, 3]));
|
|
}
|
|
}
|