From 28fe830de802eef3dcc9c3b28018ce353a4a9582 Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Tue, 10 Feb 2026 09:56:51 +0800 Subject: [PATCH 1/5] chore: initialize 0.3.0 development branch From 8a0c919925eade809e35f57a73dc4285fae7847b Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Tue, 10 Feb 2026 14:56:09 +0800 Subject: [PATCH 2/5] refactor(tensor): extract Shape and Strides into dedicated types MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce dedicated Shape and Strides types with SmallVec backing to reduce heap allocations for common tensor ranks (≤4 dimensions). This improves type safety, reduces Layout complexity, and provides a clearer separation of concerns between shape metadata and memory layout. Changes: - Add Shape type with stack allocation for common cases - Add Strides type with signed offsets for negative strides - Update Layout to use new types - Expose Shape/Strides in public API - Make TensorId crate-internal visibility --- src/autograd/mod.rs | 1 + src/lib.rs | 4 +- src/tensor/core.rs | 26 ++++++--- src/tensor/layout.rs | 130 ++++++++++++++---------------------------- src/tensor/mod.rs | 8 ++- src/tensor/shape.rs | 125 ++++++++++++++++++++++++++++++++++++++++ src/tensor/strides.rs | 123 +++++++++++++++++++++++++++++++++++++++ 7 files changed, 319 insertions(+), 98 deletions(-) create mode 100644 src/tensor/shape.rs create mode 100644 src/tensor/strides.rs diff --git a/src/autograd/mod.rs b/src/autograd/mod.rs index a48f7254..cb66a42c 100644 --- a/src/autograd/mod.rs +++ b/src/autograd/mod.rs @@ -121,6 +121,7 @@ mod forward; pub mod ops; // Reverse-mode exports +pub use crate::tensor::id::TensorId; pub use backward::{backward, backward_with_graph}; pub use grad_fn::GradFn; pub use grad_store::GradStore; diff --git a/src/lib.rs b/src/lib.rs index 91219faa..c51576d8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -96,7 +96,7 @@ pub mod prelude { // Core types pub use crate::dtype::DType; pub use crate::error::{Error, Result}; - pub use crate::tensor::{Layout, Tensor}; + pub use crate::tensor::{Layout, Shape, Strides, Tensor}; // Runtime traits pub use crate::runtime::{Device, Runtime, RuntimeClient}; @@ -116,7 +116,7 @@ pub mod prelude { // Backend runtimes #[cfg(feature = "cpu")] - pub use crate::runtime::cpu::{CpuClient, CpuDevice, CpuRuntime}; + pub use crate::runtime::cpu::{CpuClient, CpuDevice, CpuRuntime, ParallelismConfig}; #[cfg(feature = "cuda")] pub use crate::runtime::cuda::{CudaClient, CudaDevice, CudaRuntime}; diff --git a/src/tensor/core.rs b/src/tensor/core.rs index b8cc9a40..8add1c1a 100644 --- a/src/tensor/core.rs +++ b/src/tensor/core.rs @@ -63,8 +63,10 @@ impl Tensor { /// let tensor = Tensor::::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2], &device); /// # Ok::<(), numr::error::Error>(()) /// ``` + #[track_caller] pub fn from_slice(data: &[T], shape: &[usize], device: &R::Device) -> Self { - Self::try_from_slice(data, shape, device).expect("Tensor::from_slice failed") + Self::try_from_slice(data, shape, device) + .unwrap_or_else(|e| panic!("Tensor::from_slice failed: {e}")) } /// Create a tensor from a slice of data (fallible version) @@ -107,8 +109,13 @@ impl Tensor { /// /// # Safety /// The contents are uninitialized. Reading before writing is undefined behavior. + /// + /// # Panics + /// Panics if allocation fails. Use [`Self::try_empty`] in fallible contexts. + #[track_caller] pub fn empty(shape: &[usize], dtype: DType, device: &R::Device) -> Self { - Self::try_empty(shape, dtype, device).expect("Tensor::empty failed") + Self::try_empty(shape, dtype, device) + .unwrap_or_else(|e| panic!("Tensor::empty failed: {e}")) } /// Create an uninitialized tensor (fallible version) @@ -127,8 +134,10 @@ impl Tensor { /// Create a tensor filled with zeros /// /// This properly initializes memory to zero on all backends (CPU and GPU). + #[track_caller] pub fn zeros(shape: &[usize], dtype: DType, device: &R::Device) -> Self { - Self::try_zeros(shape, dtype, device).expect("Tensor::zeros failed") + Self::try_zeros(shape, dtype, device) + .unwrap_or_else(|e| panic!("Tensor::zeros failed: {e}")) } /// Create a tensor filled with zeros (fallible version) @@ -137,8 +146,9 @@ impl Tensor { } /// Create a tensor filled with ones + #[track_caller] pub fn ones(shape: &[usize], dtype: DType, device: &R::Device) -> Self { - Self::try_ones(shape, dtype, device).expect("Tensor::ones failed") + Self::try_ones(shape, dtype, device).unwrap_or_else(|e| panic!("Tensor::ones failed: {e}")) } /// Create a tensor filled with ones (fallible version) @@ -149,8 +159,10 @@ impl Tensor { /// Create a tensor filled with a scalar value /// /// The scalar is converted to the target dtype. + #[track_caller] pub fn full_scalar(shape: &[usize], dtype: DType, value: f64, device: &R::Device) -> Self { - Self::try_full_scalar(shape, dtype, value, device).expect("Tensor::full_scalar failed") + Self::try_full_scalar(shape, dtype, value, device) + .unwrap_or_else(|e| panic!("Tensor::full_scalar failed: {e}")) } /// Create a tensor filled with a scalar value (fallible version) @@ -238,9 +250,9 @@ impl Tensor { // ===== Accessors ===== - /// Get the tensor ID + /// Get the internal tensor ID for autograd graph tracking. #[inline] - pub fn id(&self) -> TensorId { + pub(crate) fn id(&self) -> TensorId { self.id } diff --git a/src/tensor/layout.rs b/src/tensor/layout.rs index 4cfc4578..02bfd1b3 100644 --- a/src/tensor/layout.rs +++ b/src/tensor/layout.rs @@ -1,19 +1,9 @@ //! Layout: shape, strides, and offset for tensor memory layout -use smallvec::SmallVec; -use std::fmt; - -/// Stack allocation threshold for dimensions -/// Most tensors have 4 or fewer dimensions, so we stack-allocate up to 4 -const STACK_DIMS: usize = 4; +pub use super::shape::Shape; +pub use super::strides::Strides; -/// Shape type: dimensions of a tensor -pub type Shape = SmallVec<[usize; STACK_DIMS]>; - -/// Strides type: element offsets between consecutive elements along each dimension -/// Signed to support negative strides (e.g., for flip operations) -/// NOTE: Strides are in ELEMENTS, not bytes (per TDD §3.0.1) -pub type Strides = SmallVec<[isize; STACK_DIMS]>; +use std::fmt; /// Layout describes the memory layout of a tensor /// @@ -54,7 +44,9 @@ impl Layout { } /// Create a layout with explicit shape, strides, and offset - pub fn new(shape: Shape, strides: Strides, offset: usize) -> Self { + pub fn new(shape: impl Into, strides: impl Into, offset: usize) -> Self { + let shape = shape.into(); + let strides = strides.into(); debug_assert_eq!(shape.len(), strides.len()); Self { shape, @@ -66,8 +58,8 @@ impl Layout { /// Create a scalar (0-dimensional) layout pub fn scalar() -> Self { Self { - shape: SmallVec::new(), - strides: SmallVec::new(), + shape: Shape::new(), + strides: Strides::new(), offset: 0, } } @@ -75,10 +67,10 @@ impl Layout { /// Compute contiguous strides for a given shape (row-major order) fn compute_contiguous_strides(shape: &[usize]) -> Strides { if shape.is_empty() { - return SmallVec::new(); + return Strides::new(); } - let mut strides: Strides = SmallVec::with_capacity(shape.len()); + let mut strides = Strides::with_capacity(shape.len()); let mut stride = 1isize; // Compute strides from last dimension to first @@ -205,12 +197,10 @@ impl Layout { /// /// Returns None if the tensor is not contiguous or shapes don't match pub fn reshape(&self, new_shape: &[usize]) -> Option { - // Must be contiguous to reshape without copying if !self.is_contiguous() { return None; } - // Element count must match let new_count: usize = new_shape.iter().product(); if new_count != self.elem_count() { return None; @@ -262,11 +252,9 @@ impl Layout { let mut new_shape = self.shape.clone(); let mut new_strides = self.strides.clone(); - // Stride for the new dimension: product of strides after this position let new_stride = if idx < ndim { new_strides[idx] * new_shape[idx] as isize } else { - // Last dimension or scalar case: stride = 1 1 }; @@ -295,12 +283,10 @@ impl Layout { pub fn permute(&self, dims: &[usize]) -> Option { let ndim = self.ndim(); - // dims must have same length as number of dimensions if dims.len() != ndim { return None; } - // Validate dims is a valid permutation: each index 0..ndim appears exactly once let mut seen = vec![false; ndim]; for &d in dims { if d >= ndim || seen[d] { @@ -309,7 +295,6 @@ impl Layout { seen[d] = true; } - // Create new shape and strides by reordering let mut new_shape = Shape::with_capacity(ndim); let mut new_strides = Strides::with_capacity(ndim); @@ -353,14 +338,11 @@ impl Layout { return None; } - // New shape: same as original but with dim narrowed let mut new_shape = self.shape.clone(); new_shape[dim] = length; - // Strides remain the same let new_strides = self.strides.clone(); - // Offset increases by start * stride[dim] let new_offset = self.offset as isize + start as isize * self.strides[dim]; if new_offset < 0 { return None; @@ -392,7 +374,6 @@ impl Layout { pub fn flip(&self, dim: isize) -> Option { let idx = self.normalize_dim(dim)?; - // If dimension size is 0 or 1, flip is a no-op if self.shape[idx] <= 1 { return Some(self.clone()); } @@ -400,17 +381,12 @@ impl Layout { let mut new_strides = self.strides.clone(); let old_stride = self.strides[idx]; - // Negate the stride new_strides[idx] = -old_stride; - // Adjust offset to point to the "last" element along this dimension - // New offset = old_offset + (dim_size - 1) * old_stride - // Use checked arithmetic to prevent overflow let dim_size = self.shape[idx] as isize; let stride_factor = (dim_size - 1).checked_mul(old_stride)?; let new_offset = (self.offset as isize).checked_add(stride_factor)?; - // Sanity check: offset should be non-negative if new_offset < 0 { return None; } @@ -450,14 +426,12 @@ impl Layout { let mut new_shape = Shape::new(); let mut new_strides = Strides::new(); - // Pad with leading 1s let pad = target.len() - self.ndim(); for &t in &target[..pad] { new_shape.push(t); - new_strides.push(0); // Stride 0 for broadcast dimensions + new_strides.push(0); } - // Check compatibility and compute strides for ((&s, &st), &t) in self .shape .iter() @@ -469,9 +443,9 @@ impl Layout { new_strides.push(st); } else if s == 1 { new_shape.push(t); - new_strides.push(0); // Broadcast: stride 0 + new_strides.push(0); } else { - return None; // Incompatible shapes + return None; } } @@ -521,6 +495,24 @@ mod tests { assert!(layout.is_contiguous()); } + #[test] + fn test_shape_and_strides_newtypes() { + let mut shape = Shape::from([2, 3, 4]); + assert_eq!(shape.len(), 3); + assert_eq!(shape.ndim(), 3); + assert!(!shape.is_empty()); + assert_eq!(shape.as_ref(), &[2, 3, 4]); + shape.remove(1); + assert_eq!(shape.as_slice(), &[2, 4]); + + let mut strides = Strides::from([12, 4, 1]); + assert_eq!(strides.len(), 3); + assert!(!strides.is_empty()); + assert_eq!(strides.as_ref(), &[12, 4, 1]); + strides.reverse(); + assert_eq!(strides.as_slice(), &[1, 4, 12]); + } + #[test] fn test_transpose() { let layout = Layout::contiguous(&[2, 3, 4]); @@ -552,8 +544,6 @@ mod tests { assert_eq!(unsqueezed.shape(), &[1, 3, 4]); } - // Note: broadcast_shape tests are in ops/arithmetic.rs - #[test] fn test_index() { let layout = Layout::contiguous(&[2, 3]); @@ -561,94 +551,69 @@ mod tests { assert_eq!(layout.index(&[0, 2]), Some(2)); assert_eq!(layout.index(&[1, 0]), Some(3)); assert_eq!(layout.index(&[1, 2]), Some(5)); - assert_eq!(layout.index(&[2, 0]), None); // Out of bounds + assert_eq!(layout.index(&[2, 0]), None); } #[test] fn test_permute() { let layout = Layout::contiguous(&[2, 3, 4]); - // Original strides: [12, 4, 1] - // Permute to [4, 2, 3] -> dims [2, 0, 1] let permuted = layout.permute(&[2, 0, 1]).unwrap(); assert_eq!(permuted.shape(), &[4, 2, 3]); assert_eq!(permuted.strides(), &[1, 12, 4]); assert!(!permuted.is_contiguous()); - // Identity permutation should preserve layout let identity = layout.permute(&[0, 1, 2]).unwrap(); assert_eq!(identity.shape(), &[2, 3, 4]); - assert_eq!(identity.strides(), &[12, 4, 1]); assert!(identity.is_contiguous()); - // Invalid permutation: wrong length assert!(layout.permute(&[0, 1]).is_none()); - - // Invalid permutation: duplicate assert!(layout.permute(&[0, 0, 1]).is_none()); - - // Invalid permutation: out of range assert!(layout.permute(&[0, 1, 5]).is_none()); } #[test] fn test_narrow() { let layout = Layout::contiguous(&[4, 5, 6]); - // Original strides: [30, 6, 1] - // Narrow dim 1: take elements 1..4 (3 elements) let narrowed = layout.narrow(1, 1, 3).unwrap(); assert_eq!(narrowed.shape(), &[4, 3, 6]); assert_eq!(narrowed.strides(), &[30, 6, 1]); - assert_eq!(narrowed.offset(), 6); // 1 * stride[1] = 1 * 6 = 6 + assert_eq!(narrowed.offset(), 6); - // Narrow dim 0: take elements 2..4 (2 elements) let narrowed2 = layout.narrow(0, 2, 2).unwrap(); assert_eq!(narrowed2.shape(), &[2, 5, 6]); - assert_eq!(narrowed2.offset(), 60); // 2 * stride[0] = 2 * 30 = 60 + assert_eq!(narrowed2.offset(), 60); - // Narrow last dim let narrowed3 = layout.narrow(2, 0, 3).unwrap(); assert_eq!(narrowed3.shape(), &[4, 5, 3]); assert_eq!(narrowed3.offset(), 0); - // Invalid: dim out of range assert!(layout.narrow(3, 0, 1).is_none()); - - // Invalid: start out of range assert!(layout.narrow(0, 5, 1).is_none()); - - // Invalid: start + length out of range assert!(layout.narrow(0, 3, 3).is_none()); - - // Invalid: length = 0 assert!(layout.narrow(0, 0, 0).is_none()); } #[test] fn test_flip() { let layout = Layout::contiguous(&[2, 3]); - // Original: strides [3, 1], offset 0 - // Flip last dim let flipped = layout.flip(-1).unwrap(); assert_eq!(flipped.shape(), &[2, 3]); - assert_eq!(flipped.strides(), &[3, -1]); // Negated stride - assert_eq!(flipped.offset(), 2); // Point to last element of row (0 + 2*1) + assert_eq!(flipped.strides(), &[3, -1]); + assert_eq!(flipped.offset(), 2); - // Flip first dim let flipped2 = layout.flip(0).unwrap(); assert_eq!(flipped2.shape(), &[2, 3]); - assert_eq!(flipped2.strides(), &[-3, 1]); // Negated stride for dim 0 - assert_eq!(flipped2.offset(), 3); // Point to last row (0 + 1*3) + assert_eq!(flipped2.strides(), &[-3, 1]); + assert_eq!(flipped2.offset(), 3); - // Flip dimension with size 1 - should be a no-op let layout1d = Layout::contiguous(&[1, 5]); let flipped1 = layout1d.flip(0).unwrap(); - assert_eq!(flipped1.strides(), &[5, 1]); // Unchanged since size is 1 + assert_eq!(flipped1.strides(), &[5, 1]); assert_eq!(flipped1.offset(), 0); - // Invalid dimension assert!(layout.flip(5).is_none()); assert!(layout.flip(-5).is_none()); } @@ -656,15 +621,11 @@ mod tests { #[test] fn test_flip_dims() { let layout = Layout::contiguous(&[2, 3, 4]); - // Original: strides [12, 4, 1], offset 0 - // Flip multiple dims let flipped = layout.flip_dims(&[0, 2]).unwrap(); assert_eq!(flipped.strides(), &[-12, 4, -1]); - // Offset: dim 0 adds (2-1)*12 = 12, dim 2 adds (4-1)*1 = 3 → total = 15 assert_eq!(flipped.offset(), 15); - // Empty dims - no change let flipped_empty = layout.flip_dims(&[]).unwrap(); assert_eq!(flipped_empty.strides(), layout.strides()); assert_eq!(flipped_empty.offset(), layout.offset()); @@ -672,16 +633,13 @@ mod tests { #[test] fn test_flip_index() { - // Test that flipped layout indexes correctly let layout = Layout::contiguous(&[2, 3]); - // Memory: [0, 1, 2, 3, 4, 5] representing [[0,1,2], [3,4,5]] - // Flip last dim: should index as [[2,1,0], [5,4,3]] let flipped = layout.flip(-1).unwrap(); - assert_eq!(flipped.index(&[0, 0]), Some(2)); // Was index 2, now at [0,0] - assert_eq!(flipped.index(&[0, 1]), Some(1)); // Was index 1, now at [0,1] - assert_eq!(flipped.index(&[0, 2]), Some(0)); // Was index 0, now at [0,2] - assert_eq!(flipped.index(&[1, 0]), Some(5)); // Was index 5 - assert_eq!(flipped.index(&[1, 2]), Some(3)); // Was index 3 + assert_eq!(flipped.index(&[0, 0]), Some(2)); + assert_eq!(flipped.index(&[0, 1]), Some(1)); + assert_eq!(flipped.index(&[0, 2]), Some(0)); + assert_eq!(flipped.index(&[1, 0]), Some(5)); + assert_eq!(flipped.index(&[1, 2]), Some(3)); } } diff --git a/src/tensor/mod.rs b/src/tensor/mod.rs index af441ac8..662d5d86 100644 --- a/src/tensor/mod.rs +++ b/src/tensor/mod.rs @@ -4,11 +4,13 @@ //! array stored on a compute device (CPU, GPU, etc.). mod core; -mod id; +pub(crate) mod id; mod layout; +pub(crate) mod shape; mod storage; +mod strides; pub use core::Tensor; -pub use id::TensorId; -pub use layout::Layout; +pub(crate) use id::TensorId; +pub use layout::{Layout, Shape, Strides}; pub use storage::Storage; diff --git a/src/tensor/shape.rs b/src/tensor/shape.rs new file mode 100644 index 00000000..a678968c --- /dev/null +++ b/src/tensor/shape.rs @@ -0,0 +1,125 @@ +//! Shape type: dimensions of a tensor + +use smallvec::SmallVec; +use std::fmt; +use std::iter::FromIterator; +use std::ops::{Deref, DerefMut}; + +/// Stack allocation threshold for dimensions +/// Most tensors have 4 or fewer dimensions, so we stack-allocate up to 4 +pub(crate) const STACK_DIMS: usize = 4; + +/// Shape type: dimensions of a tensor +#[derive(Clone, PartialEq, Eq, Default)] +pub struct Shape(SmallVec<[usize; STACK_DIMS]>); + +impl Shape { + /// Create an empty shape. + pub fn new() -> Self { + Self(SmallVec::new()) + } + + /// Create an empty shape with capacity. + pub fn with_capacity(capacity: usize) -> Self { + Self(SmallVec::with_capacity(capacity)) + } + + /// Push a dimension. + pub fn push(&mut self, dim: usize) { + self.0.push(dim); + } + + /// Remove dimension at index. + pub fn remove(&mut self, index: usize) -> usize { + self.0.remove(index) + } + + /// Insert a dimension at index. + pub fn insert(&mut self, index: usize, value: usize) { + self.0.insert(index, value); + } + + /// Swap two dimensions. + pub fn swap(&mut self, a: usize, b: usize) { + self.0.swap(a, b); + } + + /// View shape as a slice. + pub fn as_slice(&self) -> &[usize] { + self.0.as_slice() + } + + /// Number of dimensions in this shape. + #[inline] + pub fn len(&self) -> usize { + self.0.len() + } + + /// Number of dimensions in this shape. + #[inline] + pub fn ndim(&self) -> usize { + self.0.len() + } + + /// Whether this shape has zero dimensions. + #[inline] + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } +} + +impl Deref for Shape { + type Target = [usize]; + + fn deref(&self) -> &Self::Target { + self.0.as_slice() + } +} + +impl DerefMut for Shape { + fn deref_mut(&mut self) -> &mut Self::Target { + self.0.as_mut_slice() + } +} + +impl fmt::Debug for Shape { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} + +impl AsRef<[usize]> for Shape { + fn as_ref(&self) -> &[usize] { + self.0.as_slice() + } +} + +impl From> for Shape { + fn from(value: SmallVec<[usize; STACK_DIMS]>) -> Self { + Self(value) + } +} + +impl From> for Shape { + fn from(value: Vec) -> Self { + Self(value.into_iter().collect()) + } +} + +impl From<&[usize]> for Shape { + fn from(value: &[usize]) -> Self { + Self(value.iter().copied().collect()) + } +} + +impl From<[usize; N]> for Shape { + fn from(value: [usize; N]) -> Self { + Self(value.into_iter().collect()) + } +} + +impl FromIterator for Shape { + fn from_iter>(iter: T) -> Self { + Self(iter.into_iter().collect()) + } +} diff --git a/src/tensor/strides.rs b/src/tensor/strides.rs new file mode 100644 index 00000000..9e68c5e7 --- /dev/null +++ b/src/tensor/strides.rs @@ -0,0 +1,123 @@ +//! Strides type: element offsets for tensor memory layout + +use super::shape::STACK_DIMS; +use smallvec::SmallVec; +use std::fmt; +use std::iter::FromIterator; +use std::ops::{Deref, DerefMut}; + +/// Strides type: element offsets between consecutive elements along each dimension +/// Signed to support negative strides (e.g., for flip operations) +/// NOTE: Strides are in ELEMENTS, not bytes +#[derive(Clone, PartialEq, Eq, Default)] +pub struct Strides(SmallVec<[isize; STACK_DIMS]>); + +impl Strides { + /// Create empty strides. + pub fn new() -> Self { + Self(SmallVec::new()) + } + + /// Create empty strides with capacity. + pub fn with_capacity(capacity: usize) -> Self { + Self(SmallVec::with_capacity(capacity)) + } + + /// Push a stride value. + pub fn push(&mut self, stride: isize) { + self.0.push(stride); + } + + /// Remove stride at index. + pub fn remove(&mut self, index: usize) -> isize { + self.0.remove(index) + } + + /// Insert stride at index. + pub fn insert(&mut self, index: usize, value: isize) { + self.0.insert(index, value); + } + + /// Swap two strides. + pub fn swap(&mut self, a: usize, b: usize) { + self.0.swap(a, b); + } + + /// Reverse stride order. + pub fn reverse(&mut self) { + self.0.reverse(); + } + + /// View strides as a slice. + pub fn as_slice(&self) -> &[isize] { + self.0.as_slice() + } + + /// Number of stride entries. + #[inline] + pub fn len(&self) -> usize { + self.0.len() + } + + /// Whether this stride vector is empty. + #[inline] + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } +} + +impl Deref for Strides { + type Target = [isize]; + + fn deref(&self) -> &Self::Target { + self.0.as_slice() + } +} + +impl DerefMut for Strides { + fn deref_mut(&mut self) -> &mut Self::Target { + self.0.as_mut_slice() + } +} + +impl fmt::Debug for Strides { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} + +impl AsRef<[isize]> for Strides { + fn as_ref(&self) -> &[isize] { + self.0.as_slice() + } +} + +impl From> for Strides { + fn from(value: SmallVec<[isize; STACK_DIMS]>) -> Self { + Self(value) + } +} + +impl From> for Strides { + fn from(value: Vec) -> Self { + Self(value.into_iter().collect()) + } +} + +impl From<&[isize]> for Strides { + fn from(value: &[isize]) -> Self { + Self(value.iter().copied().collect()) + } +} + +impl From<[isize; N]> for Strides { + fn from(value: [isize; N]) -> Self { + Self(value.into_iter().collect()) + } +} + +impl FromIterator for Strides { + fn from_iter>(iter: T) -> Self { + Self(iter.into_iter().collect()) + } +} From a0a5e1418dfb5bd5c8c0b1f5b404ac44b27da2a7 Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Tue, 10 Feb 2026 14:56:23 +0800 Subject: [PATCH 3/5] refactor(cpu): reorganize FFT and reduce into modular structure MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Split large monolithic files into focused modules following the target architecture pattern of one operation per file. FFT changes: - Split fft.rs (single file) → fft/ directory with mod, real, shift - Add configurable parallelism via min_batch_len parameter - Improve parallel batching efficiency with par_chunks_mut Reduce changes: - Split reduce.rs → reduce/ with common, single_dim, multi_dim, precision - Expose ParallelismConfig in public API for tuning This reduces file sizes and improves code organization without changing functionality. --- src/runtime/cpu/client.rs | 144 +++- src/runtime/cpu/fft.rs | 677 ------------------- src/runtime/cpu/fft/mod.rs | 263 +++++++ src/runtime/cpu/fft/real.rs | 293 ++++++++ src/runtime/cpu/fft/shift.rs | 249 +++++++ src/runtime/cpu/helpers/reduce.rs | 526 -------------- src/runtime/cpu/helpers/reduce/common.rs | 78 +++ src/runtime/cpu/helpers/reduce/mod.rs | 146 ++++ src/runtime/cpu/helpers/reduce/multi_dim.rs | 281 ++++++++ src/runtime/cpu/helpers/reduce/precision.rs | 344 ++++++++++ src/runtime/cpu/helpers/reduce/single_dim.rs | 223 ++++++ src/runtime/cpu/kernels/fft.rs | 20 +- src/runtime/cpu/mod.rs | 2 +- 13 files changed, 2030 insertions(+), 1216 deletions(-) delete mode 100644 src/runtime/cpu/fft.rs create mode 100644 src/runtime/cpu/fft/mod.rs create mode 100644 src/runtime/cpu/fft/real.rs create mode 100644 src/runtime/cpu/fft/shift.rs delete mode 100644 src/runtime/cpu/helpers/reduce.rs create mode 100644 src/runtime/cpu/helpers/reduce/common.rs create mode 100644 src/runtime/cpu/helpers/reduce/mod.rs create mode 100644 src/runtime/cpu/helpers/reduce/multi_dim.rs create mode 100644 src/runtime/cpu/helpers/reduce/precision.rs create mode 100644 src/runtime/cpu/helpers/reduce/single_dim.rs diff --git a/src/runtime/cpu/client.rs b/src/runtime/cpu/client.rs index 30c0165f..6b3c4b9c 100644 --- a/src/runtime/cpu/client.rs +++ b/src/runtime/cpu/client.rs @@ -4,19 +4,116 @@ use super::device::CpuDevice; use super::runtime::CpuRuntime; use crate::runtime::{DefaultAllocator, RuntimeClient}; use std::alloc::{Layout as AllocLayout, alloc_zeroed, dealloc}; +#[cfg(feature = "rayon")] +use std::sync::Arc; + +/// CPU parallelism configuration. +/// +/// This is intentionally runtime-local and opt-in: +/// - `None` keeps the existing default behavior (global Rayon pool) +/// - `Some(n)` uses a dedicated pool with `n` threads for this client +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct ParallelismConfig { + /// Maximum number of worker threads for Rayon-backed kernels. + pub max_threads: Option, + /// Optional chunk size hint for future CPU kernels. + pub chunk_size: Option, +} + +impl ParallelismConfig { + /// Create a new parallelism config. + #[must_use] + pub const fn new(max_threads: Option, chunk_size: Option) -> Self { + Self { + max_threads, + chunk_size, + } + } +} /// CPU client for operation dispatch -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct CpuClient { pub(crate) device: CpuDevice, allocator: CpuAllocator, + parallelism: ParallelismConfig, + #[cfg(feature = "rayon")] + thread_pool: Option>, +} + +impl std::fmt::Debug for CpuClient { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CpuClient") + .field("device", &self.device) + .field("parallelism", &self.parallelism) + .finish() + } } impl CpuClient { /// Create a new CPU client pub fn new(device: CpuDevice) -> Self { let allocator = create_cpu_allocator(device.clone()); - Self { device, allocator } + Self { + device, + allocator, + parallelism: ParallelismConfig::default(), + #[cfg(feature = "rayon")] + thread_pool: None, + } + } + + /// Return a cloned client with explicit CPU parallelism settings. + /// + /// This allows composition with external runtimes (e.g. solvers) without + /// oversubscribing the global Rayon pool. + #[must_use] + pub fn with_parallelism(&self, config: ParallelismConfig) -> Self { + Self { + device: self.device.clone(), + allocator: self.allocator.clone(), + parallelism: config, + #[cfg(feature = "rayon")] + thread_pool: build_thread_pool(config.max_threads), + } + } + + /// Get the current parallelism settings for this client. + #[must_use] + pub const fn parallelism(&self) -> ParallelismConfig { + self.parallelism + } + + #[cfg(feature = "rayon")] + pub(crate) fn install_parallelism(&self, f: F) -> T + where + F: FnOnce() -> T + Send, + T: Send, + { + if let Some(pool) = &self.thread_pool { + pool.install(f) + } else { + f() + } + } + + #[cfg(not(feature = "rayon"))] + pub(crate) fn install_parallelism(&self, f: F) -> T + where + F: FnOnce() -> T, + { + f() + } + + #[inline] + pub(crate) fn chunk_size_hint(&self) -> usize { + self.parallelism.chunk_size.unwrap_or(1).max(1) + } + + #[cfg(feature = "rayon")] + #[inline] + pub(crate) fn rayon_min_len(&self) -> usize { + self.chunk_size_hint() } } @@ -37,6 +134,20 @@ impl RuntimeClient for CpuClient { /// CPU-specific allocator type alias pub type CpuAllocator = DefaultAllocator; +#[cfg(feature = "rayon")] +fn build_thread_pool(max_threads: Option) -> Option> { + let threads = max_threads?; + if threads == 0 { + return None; + } + + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .ok() + .map(Arc::new) +} + /// Create a CPU allocator for the given device fn create_cpu_allocator(device: CpuDevice) -> CpuAllocator { DefaultAllocator::new( @@ -67,3 +178,32 @@ fn create_cpu_allocator(device: CpuDevice) -> CpuAllocator { }, ) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::runtime::Device; + + #[test] + fn test_with_parallelism_preserves_device_and_updates_config() { + let client = CpuClient::new(CpuDevice::new()); + let configured = client.with_parallelism(ParallelismConfig::new(Some(2), Some(512))); + + assert_eq!(configured.device.id(), client.device.id()); + assert_eq!(configured.parallelism().max_threads, Some(2)); + assert_eq!(configured.parallelism().chunk_size, Some(512)); + } + + #[cfg(feature = "rayon")] + #[test] + fn test_rayon_min_len_defaults_and_normalizes_zero() { + let client = CpuClient::new(CpuDevice::new()); + assert_eq!(client.rayon_min_len(), 1); + + let configured = client.with_parallelism(ParallelismConfig::new(Some(2), Some(64))); + assert_eq!(configured.rayon_min_len(), 64); + + let zero_chunk = client.with_parallelism(ParallelismConfig::new(Some(2), Some(0))); + assert_eq!(zero_chunk.rayon_min_len(), 1); + } +} diff --git a/src/runtime/cpu/fft.rs b/src/runtime/cpu/fft.rs deleted file mode 100644 index 58649e79..00000000 --- a/src/runtime/cpu/fft.rs +++ /dev/null @@ -1,677 +0,0 @@ -//! FFT algorithm implementation for CPU runtime -//! -//! This module implements the FftAlgorithms trait for CpuClient using -//! the Stockham autosort algorithm. - -use super::{CpuClient, CpuRuntime, kernels}; -use crate::algorithm::fft::{ - FftAlgorithms, FftDirection, FftNormalization, complex_dtype_for_real, real_dtype_for_complex, - validate_fft_complex_dtype, validate_fft_size, validate_rfft_real_dtype, -}; -use crate::dtype::{Complex64, Complex128, DType}; -use crate::error::{Error, Result}; -use crate::tensor::Tensor; - -impl FftAlgorithms for CpuClient { - fn fft( - &self, - input: &Tensor, - direction: FftDirection, - norm: FftNormalization, - ) -> Result> { - // FFT along the last dimension - let ndim = input.ndim(); - if ndim == 0 { - return Err(Error::InvalidArgument { - arg: "input", - reason: "FFT requires at least 1D input".to_string(), - }); - } - self.fft_dim(input, -1, direction, norm) - } - - fn fft_dim( - &self, - input: &Tensor, - dim: isize, - direction: FftDirection, - norm: FftNormalization, - ) -> Result> { - let dtype = input.dtype(); - validate_fft_complex_dtype(dtype, "fft")?; - - let ndim = input.ndim(); - let dim = if dim < 0 { - (ndim as isize + dim) as usize - } else { - dim as usize - }; - - if dim >= ndim { - return Err(Error::InvalidDimension { - dim: dim as isize, - ndim, - }); - } - - let n = input.shape()[dim]; - validate_fft_size(n, "fft")?; - - // If FFT is along last dimension and tensor is contiguous, we can use optimized path - if dim == ndim - 1 && input.is_contiguous() { - return self.fft_last_dim_contiguous(input, direction, norm); - } - - // For non-last dimension or non-contiguous, transpose, compute, transpose back - // First, move the target dimension to the last position - let mut perm: Vec = (0..ndim).collect(); - perm.swap(dim, ndim - 1); - - let transposed = input.permute(&perm)?; - let transposed_contig = transposed.contiguous(); - - let result = self.fft_last_dim_contiguous(&transposed_contig, direction, norm)?; - - // Transpose back - result.permute(&perm) - } - - fn rfft( - &self, - input: &Tensor, - norm: FftNormalization, - ) -> Result> { - let dtype = input.dtype(); - validate_rfft_real_dtype(dtype, "rfft")?; - - let ndim = input.ndim(); - if ndim == 0 { - return Err(Error::InvalidArgument { - arg: "input", - reason: "rfft requires at least 1D input".to_string(), - }); - } - - let n = input.shape()[ndim - 1]; - validate_fft_size(n, "rfft")?; - - let input_contig = if input.is_contiguous() { - input.clone() - } else { - input.contiguous() - }; - - let output_dtype = complex_dtype_for_real(dtype)?; - let normalize_factor = norm.factor(FftDirection::Forward, n); - - // Compute output shape: [..., N/2 + 1] - let mut out_shape = input_contig.shape().to_vec(); - out_shape[ndim - 1] = n / 2 + 1; - - let output = Tensor::::empty(&out_shape, output_dtype, &self.device); - - // Compute batch size (product of all dims except last) - let batch_size: usize = input_contig.shape()[..ndim - 1].iter().product(); - let batch_size = batch_size.max(1); - - let input_ptr = input_contig.storage().ptr(); - let output_ptr = output.storage().ptr(); - - match dtype { - DType::F32 => { - let input_slice: &[f32] = - unsafe { std::slice::from_raw_parts(input_ptr as *const f32, batch_size * n) }; - let output_slice: &mut [Complex64] = unsafe { - std::slice::from_raw_parts_mut( - output_ptr as *mut Complex64, - batch_size * (n / 2 + 1), - ) - }; - - for batch_idx in 0..batch_size { - let in_start = batch_idx * n; - let out_start = batch_idx * (n / 2 + 1); - - unsafe { - kernels::rfft_c64( - &input_slice[in_start..in_start + n], - &mut output_slice[out_start..out_start + n / 2 + 1], - normalize_factor as f32, - ); - } - } - } - DType::F64 => { - let input_slice: &[f64] = - unsafe { std::slice::from_raw_parts(input_ptr as *const f64, batch_size * n) }; - let output_slice: &mut [Complex128] = unsafe { - std::slice::from_raw_parts_mut( - output_ptr as *mut Complex128, - batch_size * (n / 2 + 1), - ) - }; - - for batch_idx in 0..batch_size { - let in_start = batch_idx * n; - let out_start = batch_idx * (n / 2 + 1); - - unsafe { - kernels::rfft_c128( - &input_slice[in_start..in_start + n], - &mut output_slice[out_start..out_start + n / 2 + 1], - normalize_factor, - ); - } - } - } - _ => unreachable!(), // validate_rfft_real_dtype ensures F32 or F64 - } - - Ok(output) - } - - fn irfft( - &self, - input: &Tensor, - n: Option, - norm: FftNormalization, - ) -> Result> { - let dtype = input.dtype(); - validate_fft_complex_dtype(dtype, "irfft")?; - - let ndim = input.ndim(); - if ndim == 0 { - return Err(Error::InvalidArgument { - arg: "input", - reason: "irfft requires at least 1D input".to_string(), - }); - } - - let input_n = input.shape()[ndim - 1]; - let output_n = n.unwrap_or(2 * (input_n - 1)); - - if output_n / 2 + 1 != input_n { - return Err(Error::InvalidArgument { - arg: "n", - reason: format!( - "For irfft with n={}, input must have size {}, got {}", - output_n, - output_n / 2 + 1, - input_n - ), - }); - } - - validate_fft_size(output_n, "irfft")?; - - let input_contig = if input.is_contiguous() { - input.clone() - } else { - input.contiguous() - }; - - let output_dtype = real_dtype_for_complex(dtype)?; - let normalize_factor = norm.factor(FftDirection::Inverse, output_n); - - // Compute output shape: [..., N] - let mut out_shape = input_contig.shape().to_vec(); - out_shape[ndim - 1] = output_n; - - let output = Tensor::::empty(&out_shape, output_dtype, &self.device); - - let batch_size: usize = input_contig.shape()[..ndim - 1].iter().product(); - let batch_size = batch_size.max(1); - - let input_ptr = input_contig.storage().ptr(); - let output_ptr = output.storage().ptr(); - - match dtype { - DType::Complex64 => { - let input_slice: &[Complex64] = unsafe { - std::slice::from_raw_parts(input_ptr as *const Complex64, batch_size * input_n) - }; - let output_slice: &mut [f32] = unsafe { - std::slice::from_raw_parts_mut(output_ptr as *mut f32, batch_size * output_n) - }; - - for batch_idx in 0..batch_size { - let in_start = batch_idx * input_n; - let out_start = batch_idx * output_n; - - unsafe { - kernels::irfft_c64( - &input_slice[in_start..in_start + input_n], - &mut output_slice[out_start..out_start + output_n], - normalize_factor as f32, - ); - } - } - } - DType::Complex128 => { - let input_slice: &[Complex128] = unsafe { - std::slice::from_raw_parts(input_ptr as *const Complex128, batch_size * input_n) - }; - let output_slice: &mut [f64] = unsafe { - std::slice::from_raw_parts_mut(output_ptr as *mut f64, batch_size * output_n) - }; - - for batch_idx in 0..batch_size { - let in_start = batch_idx * input_n; - let out_start = batch_idx * output_n; - - unsafe { - kernels::irfft_c128( - &input_slice[in_start..in_start + input_n], - &mut output_slice[out_start..out_start + output_n], - normalize_factor, - ); - } - } - } - _ => unreachable!(), - } - - Ok(output) - } - - fn fft2( - &self, - input: &Tensor, - direction: FftDirection, - norm: FftNormalization, - ) -> Result> { - let ndim = input.ndim(); - if ndim < 2 { - return Err(Error::InvalidArgument { - arg: "input", - reason: "fft2 requires at least 2D input".to_string(), - }); - } - - // FFT along last two dimensions (dim=-1, then dim=-2) - // Apply normalization to both dimensions - let result = self.fft_dim(input, -1, direction, norm)?; - self.fft_dim(&result, -2, direction, norm) - } - - fn rfft2( - &self, - input: &Tensor, - norm: FftNormalization, - ) -> Result> { - let ndim = input.ndim(); - if ndim < 2 { - return Err(Error::InvalidArgument { - arg: "input", - reason: "rfft2 requires at least 2D input".to_string(), - }); - } - - // rfft along last dimension, then fft along second-to-last - let result = self.rfft(input, norm)?; - self.fft_dim(&result, -2, FftDirection::Forward, FftNormalization::None) - } - - fn irfft2( - &self, - input: &Tensor, - s: Option<(usize, usize)>, - norm: FftNormalization, - ) -> Result> { - let ndim = input.ndim(); - if ndim < 2 { - return Err(Error::InvalidArgument { - arg: "input", - reason: "irfft2 requires at least 2D input".to_string(), - }); - } - - let input_shape = input.shape(); - let (m, n) = if let Some((m, n)) = s { - (m, n) - } else { - (input_shape[ndim - 2], 2 * (input_shape[ndim - 1] - 1)) - }; - - // ifft along second-to-last, then irfft along last - let result = self.fft_dim(input, -2, FftDirection::Inverse, FftNormalization::None)?; - - // Ensure M dimension matches if specified - if result.shape()[ndim - 2] != m { - return Err(Error::InvalidArgument { - arg: "s", - reason: format!( - "M dimension mismatch: expected {}, got {}", - m, - result.shape()[ndim - 2] - ), - }); - } - - self.irfft(&result, Some(n), norm) - } - - fn fftshift(&self, input: &Tensor) -> Result> { - let dtype = input.dtype(); - let ndim = input.ndim(); - - if ndim == 0 { - return Ok(input.clone()); - } - - let input_contig = if input.is_contiguous() { - input.clone() - } else { - input.contiguous() - }; - - let n = input_contig.shape()[ndim - 1]; - let output = Tensor::::empty(input_contig.shape(), dtype, &self.device); - - let batch_size: usize = input_contig.shape()[..ndim - 1].iter().product(); - let batch_size = batch_size.max(1); - - let input_ptr = input_contig.storage().ptr(); - let output_ptr = output.storage().ptr(); - - match dtype { - DType::Complex64 => { - let input_slice: &[Complex64] = unsafe { - std::slice::from_raw_parts(input_ptr as *const Complex64, batch_size * n) - }; - let output_slice: &mut [Complex64] = unsafe { - std::slice::from_raw_parts_mut(output_ptr as *mut Complex64, batch_size * n) - }; - - for batch_idx in 0..batch_size { - let start = batch_idx * n; - unsafe { - kernels::fftshift_c64( - &input_slice[start..start + n], - &mut output_slice[start..start + n], - ); - } - } - } - DType::Complex128 => { - let input_slice: &[Complex128] = unsafe { - std::slice::from_raw_parts(input_ptr as *const Complex128, batch_size * n) - }; - let output_slice: &mut [Complex128] = unsafe { - std::slice::from_raw_parts_mut(output_ptr as *mut Complex128, batch_size * n) - }; - - for batch_idx in 0..batch_size { - let start = batch_idx * n; - unsafe { - kernels::fftshift_c128( - &input_slice[start..start + n], - &mut output_slice[start..start + n], - ); - } - } - } - _ => { - return Err(Error::UnsupportedDType { - dtype, - op: "fftshift", - }); - } - } - - Ok(output) - } - - fn ifftshift(&self, input: &Tensor) -> Result> { - let dtype = input.dtype(); - let ndim = input.ndim(); - - if ndim == 0 { - return Ok(input.clone()); - } - - let input_contig = if input.is_contiguous() { - input.clone() - } else { - input.contiguous() - }; - - let n = input_contig.shape()[ndim - 1]; - let output = Tensor::::empty(input_contig.shape(), dtype, &self.device); - - let batch_size: usize = input_contig.shape()[..ndim - 1].iter().product(); - let batch_size = batch_size.max(1); - - let input_ptr = input_contig.storage().ptr(); - let output_ptr = output.storage().ptr(); - - match dtype { - DType::Complex64 => { - let input_slice: &[Complex64] = unsafe { - std::slice::from_raw_parts(input_ptr as *const Complex64, batch_size * n) - }; - let output_slice: &mut [Complex64] = unsafe { - std::slice::from_raw_parts_mut(output_ptr as *mut Complex64, batch_size * n) - }; - - for batch_idx in 0..batch_size { - let start = batch_idx * n; - unsafe { - kernels::ifftshift_c64( - &input_slice[start..start + n], - &mut output_slice[start..start + n], - ); - } - } - } - DType::Complex128 => { - let input_slice: &[Complex128] = unsafe { - std::slice::from_raw_parts(input_ptr as *const Complex128, batch_size * n) - }; - let output_slice: &mut [Complex128] = unsafe { - std::slice::from_raw_parts_mut(output_ptr as *mut Complex128, batch_size * n) - }; - - for batch_idx in 0..batch_size { - let start = batch_idx * n; - unsafe { - kernels::ifftshift_c128( - &input_slice[start..start + n], - &mut output_slice[start..start + n], - ); - } - } - } - _ => { - return Err(Error::UnsupportedDType { - dtype, - op: "ifftshift", - }); - } - } - - Ok(output) - } - - fn fftfreq( - &self, - n: usize, - d: f64, - dtype: DType, - device: &::Device, - ) -> Result> { - if n == 0 { - return Err(Error::InvalidArgument { - arg: "n", - reason: "n must be positive".to_string(), - }); - } - - let output = Tensor::::empty(&[n], dtype, device); - let scale = 1.0 / (d * n as f64); - let output_ptr = output.storage().ptr(); - - // Frequencies: [0, 1, ..., N/2-1, -N/2, ..., -1] / (d*N) - match dtype { - DType::F32 => { - let output_slice: &mut [f32] = - unsafe { std::slice::from_raw_parts_mut(output_ptr as *mut f32, n) }; - - #[allow(clippy::needless_range_loop, clippy::manual_div_ceil)] - for i in 0..n { - let freq = if i < (n + 1) / 2 { - i as f64 - } else { - (i as isize - n as isize) as f64 - }; - output_slice[i] = (freq * scale) as f32; - } - } - DType::F64 => { - let output_slice: &mut [f64] = - unsafe { std::slice::from_raw_parts_mut(output_ptr as *mut f64, n) }; - - #[allow(clippy::needless_range_loop, clippy::manual_div_ceil)] - for i in 0..n { - let freq = if i < (n + 1) / 2 { - i as f64 - } else { - (i as isize - n as isize) as f64 - }; - output_slice[i] = freq * scale; - } - } - _ => { - return Err(Error::UnsupportedDType { - dtype, - op: "fftfreq", - }); - } - } - - Ok(output) - } - - fn rfftfreq( - &self, - n: usize, - d: f64, - dtype: DType, - device: &::Device, - ) -> Result> { - if n == 0 { - return Err(Error::InvalidArgument { - arg: "n", - reason: "n must be positive".to_string(), - }); - } - - let output_len = n / 2 + 1; - let output = Tensor::::empty(&[output_len], dtype, device); - let scale = 1.0 / (d * n as f64); - let output_ptr = output.storage().ptr(); - - // Frequencies: [0, 1, ..., N/2] / (d*N) - match dtype { - DType::F32 => { - let output_slice: &mut [f32] = - unsafe { std::slice::from_raw_parts_mut(output_ptr as *mut f32, output_len) }; - - #[allow(clippy::needless_range_loop)] - for i in 0..output_len { - output_slice[i] = (i as f64 * scale) as f32; - } - } - DType::F64 => { - let output_slice: &mut [f64] = - unsafe { std::slice::from_raw_parts_mut(output_ptr as *mut f64, output_len) }; - - #[allow(clippy::needless_range_loop)] - for i in 0..output_len { - output_slice[i] = i as f64 * scale; - } - } - _ => { - return Err(Error::UnsupportedDType { - dtype, - op: "rfftfreq", - }); - } - } - - Ok(output) - } -} - -// ============================================================================ -// Private Helper Methods -// ============================================================================ - -impl CpuClient { - /// FFT along last dimension for contiguous tensor - fn fft_last_dim_contiguous( - &self, - input: &Tensor, - direction: FftDirection, - norm: FftNormalization, - ) -> Result> { - let dtype = input.dtype(); - let ndim = input.ndim(); - let n = input.shape()[ndim - 1]; - - let inverse = matches!(direction, FftDirection::Inverse); - let normalize_factor = norm.factor(direction, n); - - let output = Tensor::::empty(input.shape(), dtype, &self.device); - - let batch_size: usize = input.shape()[..ndim - 1].iter().product(); - let batch_size = batch_size.max(1); - - let input_ptr = input.storage().ptr(); - let output_ptr = output.storage().ptr(); - - match dtype { - DType::Complex64 => { - let input_slice: &[Complex64] = unsafe { - std::slice::from_raw_parts(input_ptr as *const Complex64, batch_size * n) - }; - let output_slice: &mut [Complex64] = unsafe { - std::slice::from_raw_parts_mut(output_ptr as *mut Complex64, batch_size * n) - }; - - unsafe { - kernels::stockham_fft_batched_c64( - input_slice, - output_slice, - n, - batch_size, - inverse, - normalize_factor as f32, - ); - } - } - DType::Complex128 => { - let input_slice: &[Complex128] = unsafe { - std::slice::from_raw_parts(input_ptr as *const Complex128, batch_size * n) - }; - let output_slice: &mut [Complex128] = unsafe { - std::slice::from_raw_parts_mut(output_ptr as *mut Complex128, batch_size * n) - }; - - unsafe { - kernels::stockham_fft_batched_c128( - input_slice, - output_slice, - n, - batch_size, - inverse, - normalize_factor, - ); - } - } - _ => unreachable!(), // validate_fft_complex_dtype ensures Complex64 or Complex128 - } - - Ok(output) - } -} diff --git a/src/runtime/cpu/fft/mod.rs b/src/runtime/cpu/fft/mod.rs new file mode 100644 index 00000000..7133eefb --- /dev/null +++ b/src/runtime/cpu/fft/mod.rs @@ -0,0 +1,263 @@ +//! FFT algorithm implementation for CPU runtime +//! +//! This module implements the FftAlgorithms trait for CpuClient using +//! the Stockham autosort algorithm. + +mod real; +mod shift; + +use super::{CpuClient, CpuRuntime, kernels}; +use crate::algorithm::fft::{ + FftAlgorithms, FftDirection, FftNormalization, validate_fft_complex_dtype, validate_fft_size, +}; +use crate::dtype::{Complex64, Complex128, DType}; +use crate::error::{Error, Result}; +use crate::tensor::Tensor; + +impl FftAlgorithms for CpuClient { + fn fft( + &self, + input: &Tensor, + direction: FftDirection, + norm: FftNormalization, + ) -> Result> { + let ndim = input.ndim(); + if ndim == 0 { + return Err(Error::InvalidArgument { + arg: "input", + reason: "FFT requires at least 1D input".to_string(), + }); + } + self.fft_dim(input, -1, direction, norm) + } + + fn fft_dim( + &self, + input: &Tensor, + dim: isize, + direction: FftDirection, + norm: FftNormalization, + ) -> Result> { + let dtype = input.dtype(); + validate_fft_complex_dtype(dtype, "fft")?; + + let ndim = input.ndim(); + let dim = if dim < 0 { + (ndim as isize + dim) as usize + } else { + dim as usize + }; + + if dim >= ndim { + return Err(Error::InvalidDimension { + dim: dim as isize, + ndim, + }); + } + + let n = input.shape()[dim]; + validate_fft_size(n, "fft")?; + + if dim == ndim - 1 && input.is_contiguous() { + return self.fft_last_dim_contiguous(input, direction, norm); + } + + let mut perm: Vec = (0..ndim).collect(); + perm.swap(dim, ndim - 1); + + let transposed = input.permute(&perm)?; + let transposed_contig = transposed.contiguous(); + + let result = self.fft_last_dim_contiguous(&transposed_contig, direction, norm)?; + result.permute(&perm) + } + + fn rfft( + &self, + input: &Tensor, + norm: FftNormalization, + ) -> Result> { + real::rfft_impl(self, input, norm) + } + + fn irfft( + &self, + input: &Tensor, + n: Option, + norm: FftNormalization, + ) -> Result> { + real::irfft_impl(self, input, n, norm) + } + + fn fft2( + &self, + input: &Tensor, + direction: FftDirection, + norm: FftNormalization, + ) -> Result> { + let ndim = input.ndim(); + if ndim < 2 { + return Err(Error::InvalidArgument { + arg: "input", + reason: "fft2 requires at least 2D input".to_string(), + }); + } + + let result = self.fft_dim(input, -1, direction, norm)?; + self.fft_dim(&result, -2, direction, norm) + } + + fn rfft2( + &self, + input: &Tensor, + norm: FftNormalization, + ) -> Result> { + let ndim = input.ndim(); + if ndim < 2 { + return Err(Error::InvalidArgument { + arg: "input", + reason: "rfft2 requires at least 2D input".to_string(), + }); + } + + let result = self.rfft(input, norm)?; + self.fft_dim(&result, -2, FftDirection::Forward, FftNormalization::None) + } + + fn irfft2( + &self, + input: &Tensor, + s: Option<(usize, usize)>, + norm: FftNormalization, + ) -> Result> { + let ndim = input.ndim(); + if ndim < 2 { + return Err(Error::InvalidArgument { + arg: "input", + reason: "irfft2 requires at least 2D input".to_string(), + }); + } + + let input_shape = input.shape(); + let (m, n) = if let Some((m, n)) = s { + (m, n) + } else { + (input_shape[ndim - 2], 2 * (input_shape[ndim - 1] - 1)) + }; + + let result = self.fft_dim(input, -2, FftDirection::Inverse, FftNormalization::None)?; + + if result.shape()[ndim - 2] != m { + return Err(Error::InvalidArgument { + arg: "s", + reason: format!( + "M dimension mismatch: expected {}, got {}", + m, + result.shape()[ndim - 2] + ), + }); + } + + self.irfft(&result, Some(n), norm) + } + + fn fftshift(&self, input: &Tensor) -> Result> { + shift::fftshift_impl(self, input) + } + + fn ifftshift(&self, input: &Tensor) -> Result> { + shift::ifftshift_impl(self, input) + } + + fn fftfreq( + &self, + n: usize, + d: f64, + dtype: DType, + device: &::Device, + ) -> Result> { + shift::fftfreq_impl(n, d, dtype, device) + } + + fn rfftfreq( + &self, + n: usize, + d: f64, + dtype: DType, + device: &::Device, + ) -> Result> { + shift::rfftfreq_impl(n, d, dtype, device) + } +} + +impl CpuClient { + /// FFT along last dimension for contiguous tensor + fn fft_last_dim_contiguous( + &self, + input: &Tensor, + direction: FftDirection, + norm: FftNormalization, + ) -> Result> { + let dtype = input.dtype(); + let ndim = input.ndim(); + let n = input.shape()[ndim - 1]; + + let inverse = matches!(direction, FftDirection::Inverse); + let normalize_factor = norm.factor(direction, n); + + let output = Tensor::::empty(input.shape(), dtype, &self.device); + + let batch_size: usize = input.shape()[..ndim - 1].iter().product(); + let batch_size = batch_size.max(1); + let min_len = self.chunk_size_hint(); + + let input_ptr = input.storage().ptr(); + let output_ptr = output.storage().ptr(); + + match dtype { + DType::Complex64 => { + let input_slice: &[Complex64] = unsafe { + std::slice::from_raw_parts(input_ptr as *const Complex64, batch_size * n) + }; + let output_slice: &mut [Complex64] = unsafe { + std::slice::from_raw_parts_mut(output_ptr as *mut Complex64, batch_size * n) + }; + + self.install_parallelism(|| unsafe { + kernels::stockham_fft_batched_c64( + input_slice, + output_slice, + n, + batch_size, + inverse, + normalize_factor as f32, + min_len, + ); + }); + } + DType::Complex128 => { + let input_slice: &[Complex128] = unsafe { + std::slice::from_raw_parts(input_ptr as *const Complex128, batch_size * n) + }; + let output_slice: &mut [Complex128] = unsafe { + std::slice::from_raw_parts_mut(output_ptr as *mut Complex128, batch_size * n) + }; + + self.install_parallelism(|| unsafe { + kernels::stockham_fft_batched_c128( + input_slice, + output_slice, + n, + batch_size, + inverse, + normalize_factor, + min_len, + ); + }); + } + _ => unreachable!(), + } + + Ok(output) + } +} diff --git a/src/runtime/cpu/fft/real.rs b/src/runtime/cpu/fft/real.rs new file mode 100644 index 00000000..854423e0 --- /dev/null +++ b/src/runtime/cpu/fft/real.rs @@ -0,0 +1,293 @@ +//! Real FFT helpers (rfft, irfft) + +use super::super::{CpuClient, CpuRuntime, kernels}; +use crate::algorithm::fft::{ + FftDirection, FftNormalization, complex_dtype_for_real, real_dtype_for_complex, + validate_fft_complex_dtype, validate_fft_size, validate_rfft_real_dtype, +}; +use crate::dtype::{Complex64, Complex128, DType}; +use crate::error::{Error, Result}; +use crate::tensor::Tensor; +#[cfg(feature = "rayon")] +use rayon::prelude::*; + +pub(super) fn rfft_impl( + client: &CpuClient, + input: &Tensor, + norm: FftNormalization, +) -> Result> { + let dtype = input.dtype(); + validate_rfft_real_dtype(dtype, "rfft")?; + + let ndim = input.ndim(); + if ndim == 0 { + return Err(Error::InvalidArgument { + arg: "input", + reason: "rfft requires at least 1D input".to_string(), + }); + } + + let n = input.shape()[ndim - 1]; + validate_fft_size(n, "rfft")?; + + let input_contig = if input.is_contiguous() { + input.clone() + } else { + input.contiguous() + }; + + let output_dtype = complex_dtype_for_real(dtype)?; + let normalize_factor = norm.factor(FftDirection::Forward, n); + + let mut out_shape = input_contig.shape().to_vec(); + out_shape[ndim - 1] = n / 2 + 1; + + let output = Tensor::::empty(&out_shape, output_dtype, &client.device); + + let batch_size: usize = input_contig.shape()[..ndim - 1].iter().product(); + let batch_size = batch_size.max(1); + #[cfg(feature = "rayon")] + let min_len = client.rayon_min_len(); + + let input_ptr = input_contig.storage().ptr(); + let output_ptr = output.storage().ptr(); + + match dtype { + DType::F32 => { + let input_slice: &[f32] = + unsafe { std::slice::from_raw_parts(input_ptr as *const f32, batch_size * n) }; + let output_slice: &mut [Complex64] = unsafe { + std::slice::from_raw_parts_mut( + output_ptr as *mut Complex64, + batch_size * (n / 2 + 1), + ) + }; + let norm_f32 = normalize_factor as f32; + + client.install_parallelism(|| { + #[cfg(feature = "rayon")] + if batch_size > 1 { + output_slice + .par_chunks_mut(n / 2 + 1) + .enumerate() + .with_min_len(min_len) + .for_each(|(batch_idx, out_chunk)| { + let in_start = batch_idx * n; + unsafe { + kernels::rfft_c64( + &input_slice[in_start..in_start + n], + out_chunk, + norm_f32, + ); + } + }); + return; + } + + for batch_idx in 0..batch_size { + let in_start = batch_idx * n; + let out_start = batch_idx * (n / 2 + 1); + unsafe { + kernels::rfft_c64( + &input_slice[in_start..in_start + n], + &mut output_slice[out_start..out_start + n / 2 + 1], + norm_f32, + ); + } + } + }); + } + DType::F64 => { + let input_slice: &[f64] = + unsafe { std::slice::from_raw_parts(input_ptr as *const f64, batch_size * n) }; + let output_slice: &mut [Complex128] = unsafe { + std::slice::from_raw_parts_mut( + output_ptr as *mut Complex128, + batch_size * (n / 2 + 1), + ) + }; + + client.install_parallelism(|| { + #[cfg(feature = "rayon")] + if batch_size > 1 { + output_slice + .par_chunks_mut(n / 2 + 1) + .enumerate() + .with_min_len(min_len) + .for_each(|(batch_idx, out_chunk)| { + let in_start = batch_idx * n; + unsafe { + kernels::rfft_c128( + &input_slice[in_start..in_start + n], + out_chunk, + normalize_factor, + ); + } + }); + return; + } + + for batch_idx in 0..batch_size { + let in_start = batch_idx * n; + let out_start = batch_idx * (n / 2 + 1); + unsafe { + kernels::rfft_c128( + &input_slice[in_start..in_start + n], + &mut output_slice[out_start..out_start + n / 2 + 1], + normalize_factor, + ); + } + } + }); + } + _ => unreachable!(), + } + + Ok(output) +} + +pub(super) fn irfft_impl( + client: &CpuClient, + input: &Tensor, + n: Option, + norm: FftNormalization, +) -> Result> { + let dtype = input.dtype(); + validate_fft_complex_dtype(dtype, "irfft")?; + + let ndim = input.ndim(); + if ndim == 0 { + return Err(Error::InvalidArgument { + arg: "input", + reason: "irfft requires at least 1D input".to_string(), + }); + } + + let input_n = input.shape()[ndim - 1]; + let output_n = n.unwrap_or(2 * (input_n - 1)); + + if output_n / 2 + 1 != input_n { + return Err(Error::InvalidArgument { + arg: "n", + reason: format!( + "For irfft with n={}, input must have size {}, got {}", + output_n, + output_n / 2 + 1, + input_n + ), + }); + } + + validate_fft_size(output_n, "irfft")?; + + let input_contig = if input.is_contiguous() { + input.clone() + } else { + input.contiguous() + }; + + let output_dtype = real_dtype_for_complex(dtype)?; + let normalize_factor = norm.factor(FftDirection::Inverse, output_n); + + let mut out_shape = input_contig.shape().to_vec(); + out_shape[ndim - 1] = output_n; + + let output = Tensor::::empty(&out_shape, output_dtype, &client.device); + + let batch_size: usize = input_contig.shape()[..ndim - 1].iter().product(); + let batch_size = batch_size.max(1); + #[cfg(feature = "rayon")] + let min_len = client.rayon_min_len(); + + let input_ptr = input_contig.storage().ptr(); + let output_ptr = output.storage().ptr(); + + match dtype { + DType::Complex64 => { + let input_slice: &[Complex64] = unsafe { + std::slice::from_raw_parts(input_ptr as *const Complex64, batch_size * input_n) + }; + let output_slice: &mut [f32] = unsafe { + std::slice::from_raw_parts_mut(output_ptr as *mut f32, batch_size * output_n) + }; + let norm_f32 = normalize_factor as f32; + + client.install_parallelism(|| { + #[cfg(feature = "rayon")] + if batch_size > 1 { + output_slice + .par_chunks_mut(output_n) + .enumerate() + .with_min_len(min_len) + .for_each(|(batch_idx, out_chunk)| { + let in_start = batch_idx * input_n; + unsafe { + kernels::irfft_c64( + &input_slice[in_start..in_start + input_n], + out_chunk, + norm_f32, + ); + } + }); + return; + } + + for batch_idx in 0..batch_size { + let in_start = batch_idx * input_n; + let out_start = batch_idx * output_n; + unsafe { + kernels::irfft_c64( + &input_slice[in_start..in_start + input_n], + &mut output_slice[out_start..out_start + output_n], + norm_f32, + ); + } + } + }); + } + DType::Complex128 => { + let input_slice: &[Complex128] = unsafe { + std::slice::from_raw_parts(input_ptr as *const Complex128, batch_size * input_n) + }; + let output_slice: &mut [f64] = unsafe { + std::slice::from_raw_parts_mut(output_ptr as *mut f64, batch_size * output_n) + }; + + client.install_parallelism(|| { + #[cfg(feature = "rayon")] + if batch_size > 1 { + output_slice + .par_chunks_mut(output_n) + .enumerate() + .with_min_len(min_len) + .for_each(|(batch_idx, out_chunk)| { + let in_start = batch_idx * input_n; + unsafe { + kernels::irfft_c128( + &input_slice[in_start..in_start + input_n], + out_chunk, + normalize_factor, + ); + } + }); + return; + } + + for batch_idx in 0..batch_size { + let in_start = batch_idx * input_n; + let out_start = batch_idx * output_n; + unsafe { + kernels::irfft_c128( + &input_slice[in_start..in_start + input_n], + &mut output_slice[out_start..out_start + output_n], + normalize_factor, + ); + } + } + }); + } + _ => unreachable!(), + } + + Ok(output) +} diff --git a/src/runtime/cpu/fft/shift.rs b/src/runtime/cpu/fft/shift.rs new file mode 100644 index 00000000..b6cdb29f --- /dev/null +++ b/src/runtime/cpu/fft/shift.rs @@ -0,0 +1,249 @@ +//! FFT shift and frequency helpers + +use super::super::{CpuClient, CpuRuntime, kernels}; +use crate::dtype::{Complex64, Complex128, DType}; +use crate::error::{Error, Result}; +use crate::tensor::Tensor; +#[cfg(feature = "rayon")] +use rayon::prelude::*; + +pub(super) fn fftshift_impl( + client: &CpuClient, + input: &Tensor, +) -> Result> { + shift_impl(client, input, false) +} + +pub(super) fn ifftshift_impl( + client: &CpuClient, + input: &Tensor, +) -> Result> { + shift_impl(client, input, true) +} + +fn shift_impl( + client: &CpuClient, + input: &Tensor, + inverse: bool, +) -> Result> { + let dtype = input.dtype(); + let ndim = input.ndim(); + + if ndim == 0 { + return Ok(input.clone()); + } + + let input_contig = if input.is_contiguous() { + input.clone() + } else { + input.contiguous() + }; + + let n = input_contig.shape()[ndim - 1]; + let output = Tensor::::empty(input_contig.shape(), dtype, &client.device); + + let batch_size: usize = input_contig.shape()[..ndim - 1].iter().product(); + let batch_size = batch_size.max(1); + #[cfg(feature = "rayon")] + let min_len = client.rayon_min_len(); + + let input_ptr = input_contig.storage().ptr(); + let output_ptr = output.storage().ptr(); + + let op_name = if inverse { "ifftshift" } else { "fftshift" }; + + match dtype { + DType::Complex64 => { + let input_slice: &[Complex64] = unsafe { + std::slice::from_raw_parts(input_ptr as *const Complex64, batch_size * n) + }; + let output_slice: &mut [Complex64] = unsafe { + std::slice::from_raw_parts_mut(output_ptr as *mut Complex64, batch_size * n) + }; + + let kernel_fn = if inverse { + kernels::ifftshift_c64 + } else { + kernels::fftshift_c64 + }; + + client.install_parallelism(|| { + #[cfg(feature = "rayon")] + if batch_size > 1 { + output_slice + .par_chunks_mut(n) + .enumerate() + .with_min_len(min_len) + .for_each(|(batch_idx, out_chunk)| { + let start = batch_idx * n; + unsafe { + kernel_fn(&input_slice[start..start + n], out_chunk); + } + }); + return; + } + + for batch_idx in 0..batch_size { + let start = batch_idx * n; + unsafe { + kernel_fn( + &input_slice[start..start + n], + &mut output_slice[start..start + n], + ); + } + } + }); + } + DType::Complex128 => { + let input_slice: &[Complex128] = unsafe { + std::slice::from_raw_parts(input_ptr as *const Complex128, batch_size * n) + }; + let output_slice: &mut [Complex128] = unsafe { + std::slice::from_raw_parts_mut(output_ptr as *mut Complex128, batch_size * n) + }; + + let kernel_fn = if inverse { + kernels::ifftshift_c128 + } else { + kernels::fftshift_c128 + }; + + client.install_parallelism(|| { + #[cfg(feature = "rayon")] + if batch_size > 1 { + output_slice + .par_chunks_mut(n) + .enumerate() + .with_min_len(min_len) + .for_each(|(batch_idx, out_chunk)| { + let start = batch_idx * n; + unsafe { + kernel_fn(&input_slice[start..start + n], out_chunk); + } + }); + return; + } + + for batch_idx in 0..batch_size { + let start = batch_idx * n; + unsafe { + kernel_fn( + &input_slice[start..start + n], + &mut output_slice[start..start + n], + ); + } + } + }); + } + _ => { + return Err(Error::UnsupportedDType { dtype, op: op_name }); + } + } + + Ok(output) +} + +pub(super) fn fftfreq_impl( + n: usize, + d: f64, + dtype: DType, + device: &::Device, +) -> Result> { + if n == 0 { + return Err(Error::InvalidArgument { + arg: "n", + reason: "n must be positive".to_string(), + }); + } + + let output = Tensor::::empty(&[n], dtype, device); + let scale = 1.0 / (d * n as f64); + let output_ptr = output.storage().ptr(); + + match dtype { + DType::F32 => { + let output_slice: &mut [f32] = + unsafe { std::slice::from_raw_parts_mut(output_ptr as *mut f32, n) }; + + #[allow(clippy::needless_range_loop, clippy::manual_div_ceil)] + for i in 0..n { + let freq = if i < (n + 1) / 2 { + i as f64 + } else { + (i as isize - n as isize) as f64 + }; + output_slice[i] = (freq * scale) as f32; + } + } + DType::F64 => { + let output_slice: &mut [f64] = + unsafe { std::slice::from_raw_parts_mut(output_ptr as *mut f64, n) }; + + #[allow(clippy::needless_range_loop, clippy::manual_div_ceil)] + for i in 0..n { + let freq = if i < (n + 1) / 2 { + i as f64 + } else { + (i as isize - n as isize) as f64 + }; + output_slice[i] = freq * scale; + } + } + _ => { + return Err(Error::UnsupportedDType { + dtype, + op: "fftfreq", + }); + } + } + + Ok(output) +} + +pub(super) fn rfftfreq_impl( + n: usize, + d: f64, + dtype: DType, + device: &::Device, +) -> Result> { + if n == 0 { + return Err(Error::InvalidArgument { + arg: "n", + reason: "n must be positive".to_string(), + }); + } + + let output_len = n / 2 + 1; + let output = Tensor::::empty(&[output_len], dtype, device); + let scale = 1.0 / (d * n as f64); + let output_ptr = output.storage().ptr(); + + match dtype { + DType::F32 => { + let output_slice: &mut [f32] = + unsafe { std::slice::from_raw_parts_mut(output_ptr as *mut f32, output_len) }; + + #[allow(clippy::needless_range_loop)] + for i in 0..output_len { + output_slice[i] = (i as f64 * scale) as f32; + } + } + DType::F64 => { + let output_slice: &mut [f64] = + unsafe { std::slice::from_raw_parts_mut(output_ptr as *mut f64, output_len) }; + + #[allow(clippy::needless_range_loop)] + for i in 0..output_len { + output_slice[i] = i as f64 * scale; + } + } + _ => { + return Err(Error::UnsupportedDType { + dtype, + op: "rfftfreq", + }); + } + } + + Ok(output) +} diff --git a/src/runtime/cpu/helpers/reduce.rs b/src/runtime/cpu/helpers/reduce.rs deleted file mode 100644 index 9b3647b2..00000000 --- a/src/runtime/cpu/helpers/reduce.rs +++ /dev/null @@ -1,526 +0,0 @@ -//! Reduction operation helpers for CPU tensors - -use super::super::kernels; -use super::super::kernels::Accumulator; -use super::super::{CpuClient, CpuRuntime}; -use crate::dispatch_dtype; -use crate::dtype::Element; -use crate::error::{Error, Result}; -use crate::ops::{AccumulationPrecision, Kernel, ReduceOp, reduce_output_shape}; -use crate::runtime::ensure_contiguous; -use crate::tensor::Tensor; - -/// Reduce implementation with native precision -pub fn reduce_impl( - client: &CpuClient, - op: ReduceOp, - a: &Tensor, - dims: &[usize], - keepdim: bool, - op_name: &'static str, -) -> Result> { - let dtype = a.dtype(); - let shape = a.shape(); - let ndim = shape.len(); - - // Validate dimensions - for &d in dims { - if d >= ndim { - return Err(Error::InvalidDimension { - dim: d as isize, - ndim, - }); - } - } - - // Fast path: reduce last dimension when contiguous (uses SIMD kernel) - // Multi-dimension reduction uses sequential single-dimension reductions - if dims.len() == 1 && dims[0] == ndim - 1 && a.is_contiguous() { - // Simple case: reduce last dimension - let reduce_size = shape[ndim - 1]; - let outer_size: usize = shape[..ndim - 1].iter().product(); - let outer_size = outer_size.max(1); - - let out_shape = reduce_output_shape(shape, dims, keepdim); - let out = Tensor::::empty(&out_shape, dtype, &client.device); - - let a_ptr = a.storage().ptr(); - let out_ptr = out.storage().ptr(); - - dispatch_dtype!(dtype, T => { - unsafe { - >::reduce::( - client, - op, - a_ptr as *const T, - out_ptr as *mut T, - reduce_size, - outer_size, - ); - } - }, op_name); - - Ok(out) - } else if dims.is_empty() { - // No dimensions to reduce - return a copy - Ok(a.clone()) - } else { - // General case: reduce multiple/arbitrary dimensions - // Sequential reduction (highest dim first) is standard approach - // used by PyTorch and other frameworks - let a_contig = ensure_contiguous(a); - - // Reduce one dimension at a time, from highest to lowest - // (reduces index shifting issues as shape shrinks) - let mut sorted_dims: Vec = dims.to_vec(); - sorted_dims.sort_unstable(); - sorted_dims.reverse(); - - let mut current = a_contig; - for &dim in &sorted_dims { - current = reduce_single_dim(client, op, ¤t, dim, keepdim, op_name)?; - } - - Ok(current) - } -} - -/// Reduce a single dimension of a tensor using native precision. -/// -/// This is an optimized path for single-dimension reductions. Uses chunked -/// iteration for non-last dimensions to handle strided memory access. -/// -/// # Arguments -/// * `op` - Reduction operation (Sum, Max, Min) -/// * `dim` - Dimension to reduce -/// * `keepdim` - Whether to keep the reduced dimension as size 1 -/// * `op_name` - Name for error messages -fn reduce_single_dim( - client: &CpuClient, - op: ReduceOp, - a: &Tensor, - dim: usize, - keepdim: bool, - op_name: &'static str, -) -> Result> { - let dtype = a.dtype(); - let shape = a.shape(); - let ndim = shape.len(); - - if dim >= ndim { - return Err(Error::InvalidDimension { - dim: dim as isize, - ndim, - }); - } - - let reduce_size = shape[dim]; - let outer_size: usize = shape[..dim].iter().product(); - let outer_size = outer_size.max(1); - let inner_size: usize = shape[dim + 1..].iter().product(); - let inner_size = inner_size.max(1); - - let out_shape = reduce_output_shape(shape, &[dim], keepdim); - let out = Tensor::::empty(&out_shape, dtype, &client.device); - - // If reducing non-last dimension, we need special handling - if dim == ndim - 1 { - // Reducing last dimension - can use kernel directly - let a_ptr = a.storage().ptr(); - let out_ptr = out.storage().ptr(); - - dispatch_dtype!(dtype, T => { - unsafe { - >::reduce::( - client, - op, - a_ptr as *const T, - out_ptr as *mut T, - reduce_size, - outer_size, - ); - } - }, op_name); - } else { - // Reducing non-last dimension - iterate manually - let a_ptr = a.storage().ptr(); - let out_ptr = out.storage().ptr(); - - dispatch_dtype!(dtype, T => { - unsafe { - reduce_non_last_dim::( - op, - a_ptr as *const T, - out_ptr as *mut T, - outer_size, - reduce_size, - inner_size, - ); - } - }, op_name); - } - - Ok(out) -} - -/// Reduce a non-last dimension -#[allow(unsafe_op_in_unsafe_fn)] -unsafe fn reduce_non_last_dim( - op: ReduceOp, - a: *const T, - out: *mut T, - outer_size: usize, - reduce_size: usize, - inner_size: usize, -) { - for outer in 0..outer_size { - for inner in 0..inner_size { - let mut acc = match op { - ReduceOp::Sum | ReduceOp::Mean => T::zero(), - ReduceOp::Prod => T::one(), - ReduceOp::Max => { - let idx = outer * reduce_size * inner_size + inner; - *a.add(idx) - } - ReduceOp::Min => { - let idx = outer * reduce_size * inner_size + inner; - *a.add(idx) - } - ReduceOp::All => T::one(), - ReduceOp::Any => T::zero(), - }; - - for r in 0..reduce_size { - let idx = outer * reduce_size * inner_size + r * inner_size + inner; - let val = *a.add(idx); - - acc = match op { - ReduceOp::Sum | ReduceOp::Mean => acc + val, - ReduceOp::Prod => acc * val, - ReduceOp::Max => { - if val > acc { - val - } else { - acc - } - } - ReduceOp::Min => { - if val < acc { - val - } else { - acc - } - } - ReduceOp::All => { - if val.to_f64() != 0.0 && acc.to_f64() != 0.0 { - T::one() - } else { - T::zero() - } - } - ReduceOp::Any => { - if val.to_f64() != 0.0 || acc.to_f64() != 0.0 { - T::one() - } else { - T::zero() - } - } - }; - } - - // Apply mean scaling if needed - if matches!(op, ReduceOp::Mean) { - acc = T::from_f64(acc.to_f64() / reduce_size as f64); - } - - let out_idx = outer * inner_size + inner; - *out.add(out_idx) = acc; - } - } -} - -// ============================================================================ -// Precision-Aware Reduction Helpers -// ============================================================================ - -/// Reduce implementation with explicit accumulation precision -pub fn reduce_impl_with_precision( - client: &CpuClient, - op: ReduceOp, - a: &Tensor, - dims: &[usize], - keepdim: bool, - precision: AccumulationPrecision, - op_name: &'static str, -) -> Result> { - let dtype = a.dtype(); - let shape = a.shape(); - let ndim = shape.len(); - - // Validate dimensions - for &d in dims { - if d >= ndim { - return Err(Error::InvalidDimension { - dim: d as isize, - ndim, - }); - } - } - - // Fast path: reduce last dimension when contiguous (uses SIMD kernel) - if dims.len() == 1 && dims[0] == ndim - 1 && a.is_contiguous() { - // Simple case: reduce last dimension - let reduce_size = shape[ndim - 1]; - let outer_size: usize = shape[..ndim - 1].iter().product(); - let outer_size = outer_size.max(1); - - let out_shape = reduce_output_shape(shape, dims, keepdim); - let out = Tensor::::empty(&out_shape, dtype, &client.device); - - let a_ptr = a.storage().ptr(); - let out_ptr = out.storage().ptr(); - - dispatch_dtype!(dtype, T => { - unsafe { - kernels::reduce_kernel_with_precision::( - op, - a_ptr as *const T, - out_ptr as *mut T, - reduce_size, - outer_size, - precision, - ); - } - }, op_name); - - Ok(out) - } else if dims.is_empty() { - // No dimensions to reduce - return a copy - Ok(a.clone()) - } else { - // General case: reduce multiple/arbitrary dimensions - // Sequential reduction (highest dim first) is standard approach - let a_contig = ensure_contiguous(a); - - // Reduce one dimension at a time, from highest to lowest - // (reduces index shifting issues as shape shrinks) - let mut sorted_dims: Vec = dims.to_vec(); - sorted_dims.sort_unstable(); - sorted_dims.reverse(); - - let mut current = a_contig; - for &dim in &sorted_dims { - current = reduce_single_dim_with_precision( - client, op, ¤t, dim, keepdim, precision, op_name, - )?; - } - - Ok(current) - } -} - -/// Reduce a single dimension with explicit accumulation precision. -/// -/// Similar to [`reduce_single_dim`], but allows specifying higher precision -/// for accumulation. This is important for numerical stability when summing -/// many small values (e.g., F16 tensors accumulated in F32). -/// -/// # Arguments -/// * `op` - Reduction operation (Sum, Max, Min) -/// * `dim` - Dimension to reduce -/// * `keepdim` - Whether to keep the reduced dimension as size 1 -/// * `precision` - Accumulation precision (Native, F32, or F64) -/// * `op_name` - Name for error messages -fn reduce_single_dim_with_precision( - client: &CpuClient, - op: ReduceOp, - a: &Tensor, - dim: usize, - keepdim: bool, - precision: AccumulationPrecision, - op_name: &'static str, -) -> Result> { - let dtype = a.dtype(); - let shape = a.shape(); - let ndim = shape.len(); - - if dim >= ndim { - return Err(Error::InvalidDimension { - dim: dim as isize, - ndim, - }); - } - - let reduce_size = shape[dim]; - let outer_size: usize = shape[..dim].iter().product(); - let outer_size = outer_size.max(1); - let inner_size: usize = shape[dim + 1..].iter().product(); - let inner_size = inner_size.max(1); - - let out_shape = reduce_output_shape(shape, &[dim], keepdim); - let out = Tensor::::empty(&out_shape, dtype, &client.device); - - let a_ptr = a.storage().ptr(); - let out_ptr = out.storage().ptr(); - - if dim == ndim - 1 { - // Reducing last dimension - use kernel directly - dispatch_dtype!(dtype, T => { - unsafe { - kernels::reduce_kernel_with_precision::( - op, - a_ptr as *const T, - out_ptr as *mut T, - reduce_size, - outer_size, - precision, - ); - } - }, op_name); - } else { - // Reducing non-last dimension - iterate manually with precision - dispatch_dtype!(dtype, T => { - unsafe { - reduce_non_last_dim_with_precision::( - op, - a_ptr as *const T, - out_ptr as *mut T, - outer_size, - reduce_size, - inner_size, - precision, - ); - } - }, op_name); - } - - Ok(out) -} - -/// Reduce a non-last dimension with explicit precision -#[allow(unsafe_op_in_unsafe_fn)] -unsafe fn reduce_non_last_dim_with_precision( - op: ReduceOp, - a: *const T, - out: *mut T, - outer_size: usize, - reduce_size: usize, - inner_size: usize, - precision: AccumulationPrecision, -) { - match precision { - AccumulationPrecision::Native => { - reduce_non_last_dim(op, a, out, outer_size, reduce_size, inner_size); - } - AccumulationPrecision::FP32 | AccumulationPrecision::BF16 => { - // For CPU, both FP32 and BF16 accumulation use f32 internally - reduce_non_last_dim_f32_acc(op, a, out, outer_size, reduce_size, inner_size); - } - AccumulationPrecision::FP64 => { - // Maximum precision for math/science applications - reduce_non_last_dim_f64_acc(op, a, out, outer_size, reduce_size, inner_size); - } - } -} - -/// Generic reduce for non-last dimension with configurable accumulation precision. -/// -/// Uses the Accumulator trait to abstract over f32/f64 accumulation types. -#[allow(unsafe_op_in_unsafe_fn)] -unsafe fn reduce_non_last_dim_acc( - op: ReduceOp, - a: *const T, - out: *mut T, - outer_size: usize, - reduce_size: usize, - inner_size: usize, -) { - for outer in 0..outer_size { - for inner in 0..inner_size { - let first_idx = outer * reduce_size * inner_size + inner; - let first_val = A::acc_in((*a.add(first_idx)).to_f64()); - - let mut acc: A = match op { - ReduceOp::Sum | ReduceOp::Mean => A::ZERO, - ReduceOp::Prod => A::ONE, - ReduceOp::Max | ReduceOp::Min => first_val, - ReduceOp::All => A::ONE, - ReduceOp::Any => A::ZERO, - }; - - for r in 0..reduce_size { - let idx = outer * reduce_size * inner_size + r * inner_size + inner; - let val = A::acc_in((*a.add(idx)).to_f64()); - - acc = match op { - ReduceOp::Sum | ReduceOp::Mean => acc.acc_add(val), - ReduceOp::Prod => acc.acc_mul(val), - ReduceOp::Max => { - if val > acc { - val - } else { - acc - } - } - ReduceOp::Min => { - if val < acc { - val - } else { - acc - } - } - ReduceOp::All => { - if val != A::ZERO && acc != A::ZERO { - A::ONE - } else { - A::ZERO - } - } - ReduceOp::Any => { - if val != A::ZERO || acc != A::ZERO { - A::ONE - } else { - A::ZERO - } - } - }; - } - - // Apply mean scaling if needed - if matches!(op, ReduceOp::Mean) { - acc = acc.acc_div(reduce_size); - } - - let out_idx = outer * inner_size + inner; - *out.add(out_idx) = T::from_f64(acc.into()); - } - } -} - -/// Reduce a non-last dimension with f32 accumulation (convenience wrapper) -#[allow(unsafe_op_in_unsafe_fn)] -#[inline] -unsafe fn reduce_non_last_dim_f32_acc( - op: ReduceOp, - a: *const T, - out: *mut T, - outer_size: usize, - reduce_size: usize, - inner_size: usize, -) { - reduce_non_last_dim_acc::(op, a, out, outer_size, reduce_size, inner_size) -} - -/// Reduce a non-last dimension with f64 accumulation (convenience wrapper) -#[allow(unsafe_op_in_unsafe_fn)] -#[inline] -unsafe fn reduce_non_last_dim_f64_acc( - op: ReduceOp, - a: *const T, - out: *mut T, - outer_size: usize, - reduce_size: usize, - inner_size: usize, -) { - reduce_non_last_dim_acc::(op, a, out, outer_size, reduce_size, inner_size) -} diff --git a/src/runtime/cpu/helpers/reduce/common.rs b/src/runtime/cpu/helpers/reduce/common.rs new file mode 100644 index 00000000..bb934cac --- /dev/null +++ b/src/runtime/cpu/helpers/reduce/common.rs @@ -0,0 +1,78 @@ +//! Shared helpers for reduction operations + +/// Use a fused single-pass path for small multi-dimension contiguous reductions. +/// +/// This avoids intermediate tensor allocations from sequential dim-by-dim reduction. +/// The 1 MiB threshold is chosen to fit comfortably within L1/L2 cache on modern CPUs, +/// ensuring the single-pass scan has good cache locality. +pub(super) const FUSED_MULTI_DIM_REDUCTION_MAX_BYTES: usize = 1 << 20; // 1 MiB + +use crate::runtime::cpu::CpuRuntime; +use crate::tensor::Tensor; + +#[inline] +pub(super) fn should_fuse_multi_dim_reduction(a: &Tensor, dims: &[usize]) -> bool { + if dims.len() <= 1 || !a.is_contiguous() { + return false; + } + + if a.numel() == 0 { + return false; + } + + let bytes = a.numel().saturating_mul(a.dtype().size_in_bytes()); + bytes <= FUSED_MULTI_DIM_REDUCTION_MAX_BYTES +} + +#[inline] +pub(super) fn contiguous_strides(shape: &[usize]) -> Vec { + if shape.is_empty() { + return Vec::new(); + } + + let mut strides = vec![1usize; shape.len()]; + for i in (0..shape.len() - 1).rev() { + strides[i] = strides[i + 1] * shape[i + 1]; + } + strides +} + +#[inline] +pub(super) fn out_index_from_coord( + coord: &[usize], + reduce_mask: &[bool], + keepdim: bool, + kept_axes: &[usize], + out_strides: &[usize], +) -> usize { + if out_strides.is_empty() { + return 0; + } + + if keepdim { + let mut idx = 0usize; + for axis in 0..coord.len() { + if !reduce_mask[axis] { + idx += coord[axis] * out_strides[axis]; + } + } + idx + } else { + let mut idx = 0usize; + for (out_axis, &axis) in kept_axes.iter().enumerate() { + idx += coord[axis] * out_strides[out_axis]; + } + idx + } +} + +#[inline] +pub(super) fn advance_coord(coord: &mut [usize], shape: &[usize]) { + for axis in (0..coord.len()).rev() { + coord[axis] += 1; + if coord[axis] < shape[axis] { + return; + } + coord[axis] = 0; + } +} diff --git a/src/runtime/cpu/helpers/reduce/mod.rs b/src/runtime/cpu/helpers/reduce/mod.rs new file mode 100644 index 00000000..39999be4 --- /dev/null +++ b/src/runtime/cpu/helpers/reduce/mod.rs @@ -0,0 +1,146 @@ +//! Reduction operation helpers for CPU tensors + +mod common; +mod multi_dim; +mod precision; +mod single_dim; + +pub use precision::reduce_impl_with_precision; + +use common::should_fuse_multi_dim_reduction; +use multi_dim::reduce_multi_dim_fused; +use single_dim::reduce_single_dim; + +use crate::dispatch_dtype; +use crate::error::{Error, Result}; +use crate::ops::{AccumulationPrecision, Kernel, ReduceOp, reduce_output_shape}; +use crate::runtime::cpu::{CpuClient, CpuRuntime}; +use crate::runtime::ensure_contiguous; +use crate::tensor::Tensor; + +/// Reduce implementation with native precision +pub fn reduce_impl( + client: &CpuClient, + op: ReduceOp, + a: &Tensor, + dims: &[usize], + keepdim: bool, + op_name: &'static str, +) -> Result> { + let dtype = a.dtype(); + let shape = a.shape(); + let ndim = shape.len(); + + for &d in dims { + if d >= ndim { + return Err(Error::InvalidDimension { + dim: d as isize, + ndim, + }); + } + } + + // Fast path: reduce last dimension when contiguous (uses SIMD kernel) + if dims.len() == 1 && dims[0] == ndim - 1 && a.is_contiguous() { + let reduce_size = shape[ndim - 1]; + let outer_size: usize = shape[..ndim - 1].iter().product(); + let outer_size = outer_size.max(1); + + let out_shape = reduce_output_shape(shape, dims, keepdim); + let out = Tensor::::empty(&out_shape, dtype, &client.device); + + let a_ptr = a.storage().ptr(); + let out_ptr = out.storage().ptr(); + + dispatch_dtype!(dtype, T => { + unsafe { + >::reduce::( + client, + op, + a_ptr as *const T, + out_ptr as *mut T, + reduce_size, + outer_size, + ); + } + }, op_name); + + Ok(out) + } else if dims.is_empty() { + Ok(a.clone()) + } else if should_fuse_multi_dim_reduction(a, dims) { + reduce_multi_dim_fused( + client, + op, + a, + dims, + keepdim, + AccumulationPrecision::Native, + op_name, + ) + } else { + let a_contig = ensure_contiguous(a); + + let mut sorted_dims: Vec = dims.to_vec(); + sorted_dims.sort_unstable(); + sorted_dims.reverse(); + + let mut current = a_contig; + for &dim in &sorted_dims { + current = reduce_single_dim(client, op, ¤t, dim, keepdim, op_name)?; + } + + Ok(current) + } +} + +#[cfg(test)] +mod tests { + use crate::ops::{AccumulationPrecision, ReduceOps}; + use crate::runtime::Runtime; + use crate::runtime::cpu::{CpuDevice, CpuRuntime}; + use crate::tensor::Tensor; + + #[test] + fn test_fused_multi_dim_sum_matches_expected() { + let device = CpuDevice::new(); + let client = CpuRuntime::default_client(&device); + let data: Vec = (1..=24).map(|v| v as f32).collect(); + let a = Tensor::::from_slice(&data, &[2, 3, 4], &device); + + let out = client.sum(&a, &[1, 2], false).unwrap(); + let got: Vec = out.to_vec(); + assert_eq!(got, vec![78.0, 222.0]); + } + + #[test] + fn test_fused_multi_dim_mean_keepdim_matches_expected() { + let device = CpuDevice::new(); + let client = CpuRuntime::default_client(&device); + let data: Vec = (1..=24).map(|v| v as f32).collect(); + let a = Tensor::::from_slice(&data, &[2, 3, 4], &device); + + let out = client.mean(&a, &[0, 2], true).unwrap(); + assert_eq!(out.shape(), &[1, 3, 1]); + let got: Vec = out.to_vec(); + assert_eq!(got, vec![8.5, 12.5, 16.5]); + } + + #[test] + fn test_fused_multi_dim_max_and_precision_sum() { + let device = CpuDevice::new(); + let client = CpuRuntime::default_client(&device); + let data: Vec = (1..=24).map(|v| v as f32).collect(); + let a = Tensor::::from_slice(&data, &[2, 3, 4], &device); + + let max_out = client.max(&a, &[0, 1], false).unwrap(); + let max_vals: Vec = max_out.to_vec(); + assert_eq!(max_vals, vec![21.0, 22.0, 23.0, 24.0]); + + let sum_prec = client + .sum_with_precision(&a, &[0, 2], false, AccumulationPrecision::FP64) + .unwrap(); + let sum_vals: Vec = sum_prec.to_vec(); + assert_eq!(sum_vals, vec![68.0, 100.0, 132.0]); + } +} diff --git a/src/runtime/cpu/helpers/reduce/multi_dim.rs b/src/runtime/cpu/helpers/reduce/multi_dim.rs new file mode 100644 index 00000000..83ccf2cc --- /dev/null +++ b/src/runtime/cpu/helpers/reduce/multi_dim.rs @@ -0,0 +1,281 @@ +//! Fused contiguous multi-dimension reduction + +use super::common::{advance_coord, contiguous_strides, out_index_from_coord}; +use crate::dispatch_dtype; +use crate::dtype::Element; +use crate::error::Result; +use crate::ops::{AccumulationPrecision, ReduceOp, reduce_output_shape}; +use crate::runtime::cpu::kernels::Accumulator; +use crate::runtime::cpu::{CpuClient, CpuRuntime}; +use crate::tensor::Tensor; + +/// Fused contiguous multi-dimension reduction for small tensors. +/// +/// Executes reduction in a single pass over input elements and writes directly +/// into output buckets, avoiding intermediate tensors from repeated single-dim +/// reductions. +pub(super) fn reduce_multi_dim_fused( + client: &CpuClient, + op: ReduceOp, + a: &Tensor, + dims: &[usize], + keepdim: bool, + precision: AccumulationPrecision, + op_name: &'static str, +) -> Result> { + let shape = a.shape(); + let out_shape = reduce_output_shape(shape, dims, keepdim); + let out = Tensor::::empty(&out_shape, a.dtype(), &client.device); + + let mut reduce_mask = vec![false; shape.len()]; + for &d in dims { + reduce_mask[d] = true; + } + + let kept_axes: Vec = if keepdim { + Vec::new() + } else { + (0..shape.len()) + .filter(|&axis| !reduce_mask[axis]) + .collect() + }; + let out_strides = contiguous_strides(&out_shape); + let reduce_count = dims.iter().fold(1usize, |acc, &d| acc * shape[d]); + let numel = a.numel(); + let out_numel = out.numel(); + + let in_ptr = a.storage().ptr(); + let out_ptr = out.storage().ptr(); + + dispatch_dtype!(a.dtype(), T => { + unsafe { + match precision { + AccumulationPrecision::Native => reduce_multi_dim_fused_native::( + op, + in_ptr as *const T, + out_ptr as *mut T, + numel, + out_numel, + shape, + &reduce_mask, + keepdim, + &kept_axes, + &out_strides, + reduce_count, + ), + AccumulationPrecision::FP32 | AccumulationPrecision::BF16 => { + reduce_multi_dim_fused_acc::( + op, + in_ptr as *const T, + out_ptr as *mut T, + numel, + out_numel, + shape, + &reduce_mask, + keepdim, + &kept_axes, + &out_strides, + reduce_count, + ) + } + AccumulationPrecision::FP64 => reduce_multi_dim_fused_acc::( + op, + in_ptr as *const T, + out_ptr as *mut T, + numel, + out_numel, + shape, + &reduce_mask, + keepdim, + &kept_axes, + &out_strides, + reduce_count, + ), + } + } + }, op_name); + + Ok(out) +} + +#[allow(unsafe_op_in_unsafe_fn)] +unsafe fn reduce_multi_dim_fused_native( + op: ReduceOp, + input: *const T, + output: *mut T, + numel: usize, + out_numel: usize, + shape: &[usize], + reduce_mask: &[bool], + keepdim: bool, + kept_axes: &[usize], + out_strides: &[usize], + reduce_count: usize, +) { + match op { + ReduceOp::Sum | ReduceOp::Mean | ReduceOp::Any => { + for i in 0..out_numel { + *output.add(i) = T::zero(); + } + } + ReduceOp::Prod | ReduceOp::All => { + for i in 0..out_numel { + *output.add(i) = T::one(); + } + } + ReduceOp::Max | ReduceOp::Min => {} + } + + let mut initialized = if matches!(op, ReduceOp::Max | ReduceOp::Min) { + vec![false; out_numel] + } else { + Vec::new() + }; + + let mut coord = vec![0usize; shape.len()]; + for linear in 0..numel { + let out_idx = out_index_from_coord(&coord, reduce_mask, keepdim, kept_axes, out_strides); + let val = *input.add(linear); + + match op { + ReduceOp::Sum | ReduceOp::Mean => { + let acc = *output.add(out_idx); + *output.add(out_idx) = acc + val; + } + ReduceOp::Prod => { + let acc = *output.add(out_idx); + *output.add(out_idx) = acc * val; + } + ReduceOp::Max => { + if !initialized[out_idx] { + *output.add(out_idx) = val; + initialized[out_idx] = true; + } else { + let acc = *output.add(out_idx); + *output.add(out_idx) = if val > acc { val } else { acc }; + } + } + ReduceOp::Min => { + if !initialized[out_idx] { + *output.add(out_idx) = val; + initialized[out_idx] = true; + } else { + let acc = *output.add(out_idx); + *output.add(out_idx) = if val < acc { val } else { acc }; + } + } + ReduceOp::All => { + let acc = *output.add(out_idx); + *output.add(out_idx) = if val.to_f64() != 0.0 && acc.to_f64() != 0.0 { + T::one() + } else { + T::zero() + }; + } + ReduceOp::Any => { + let acc = *output.add(out_idx); + *output.add(out_idx) = if val.to_f64() != 0.0 || acc.to_f64() != 0.0 { + T::one() + } else { + T::zero() + }; + } + } + + if linear + 1 < numel { + advance_coord(&mut coord, shape); + } + } + + if matches!(op, ReduceOp::Mean) { + for i in 0..out_numel { + let scaled = (*output.add(i)).to_f64() / reduce_count as f64; + *output.add(i) = T::from_f64(scaled); + } + } +} + +#[allow(unsafe_op_in_unsafe_fn)] +unsafe fn reduce_multi_dim_fused_acc( + op: ReduceOp, + input: *const T, + output: *mut T, + numel: usize, + out_numel: usize, + shape: &[usize], + reduce_mask: &[bool], + keepdim: bool, + kept_axes: &[usize], + out_strides: &[usize], + reduce_count: usize, +) { + let mut acc = match op { + ReduceOp::Sum | ReduceOp::Mean | ReduceOp::Any | ReduceOp::Max | ReduceOp::Min => { + vec![A::ZERO; out_numel] + } + ReduceOp::Prod | ReduceOp::All => vec![A::ONE; out_numel], + }; + + let mut initialized = if matches!(op, ReduceOp::Max | ReduceOp::Min) { + vec![false; out_numel] + } else { + Vec::new() + }; + + let mut coord = vec![0usize; shape.len()]; + for linear in 0..numel { + let out_idx = out_index_from_coord(&coord, reduce_mask, keepdim, kept_axes, out_strides); + let val = A::acc_in((*input.add(linear)).to_f64()); + + match op { + ReduceOp::Sum | ReduceOp::Mean => { + acc[out_idx] = acc[out_idx].acc_add(val); + } + ReduceOp::Prod => { + acc[out_idx] = acc[out_idx].acc_mul(val); + } + ReduceOp::Max => { + if !initialized[out_idx] { + acc[out_idx] = val; + initialized[out_idx] = true; + } else if val > acc[out_idx] { + acc[out_idx] = val; + } + } + ReduceOp::Min => { + if !initialized[out_idx] { + acc[out_idx] = val; + initialized[out_idx] = true; + } else if val < acc[out_idx] { + acc[out_idx] = val; + } + } + ReduceOp::All => { + acc[out_idx] = if val != A::ZERO && acc[out_idx] != A::ZERO { + A::ONE + } else { + A::ZERO + }; + } + ReduceOp::Any => { + acc[out_idx] = if val != A::ZERO || acc[out_idx] != A::ZERO { + A::ONE + } else { + A::ZERO + }; + } + } + + if linear + 1 < numel { + advance_coord(&mut coord, shape); + } + } + + for i in 0..out_numel { + let mut out_val = acc[i]; + if matches!(op, ReduceOp::Mean) { + out_val = out_val.acc_div(reduce_count); + } + *output.add(i) = T::from_f64(out_val.into()); + } +} diff --git a/src/runtime/cpu/helpers/reduce/precision.rs b/src/runtime/cpu/helpers/reduce/precision.rs new file mode 100644 index 00000000..68f0e5ac --- /dev/null +++ b/src/runtime/cpu/helpers/reduce/precision.rs @@ -0,0 +1,344 @@ +//! Precision-aware reduction helpers (FP32/FP64 accumulation) + +use super::multi_dim::reduce_multi_dim_fused; +use super::single_dim::reduce_non_last_dim_outer; +use crate::dispatch_dtype; +use crate::dtype::Element; +use crate::error::{Error, Result}; +use crate::ops::{AccumulationPrecision, ReduceOp, reduce_output_shape}; +use crate::runtime::cpu::kernels::{self, Accumulator}; +use crate::runtime::cpu::{CpuClient, CpuRuntime}; +use crate::runtime::ensure_contiguous; +use crate::tensor::Tensor; +#[cfg(feature = "rayon")] +use rayon::prelude::*; + +use super::common::should_fuse_multi_dim_reduction; + +/// Reduce implementation with explicit accumulation precision +pub fn reduce_impl_with_precision( + client: &CpuClient, + op: ReduceOp, + a: &Tensor, + dims: &[usize], + keepdim: bool, + precision: AccumulationPrecision, + op_name: &'static str, +) -> Result> { + let dtype = a.dtype(); + let shape = a.shape(); + let ndim = shape.len(); + + for &d in dims { + if d >= ndim { + return Err(Error::InvalidDimension { + dim: d as isize, + ndim, + }); + } + } + + if dims.len() == 1 && dims[0] == ndim - 1 && a.is_contiguous() { + let reduce_size = shape[ndim - 1]; + let outer_size: usize = shape[..ndim - 1].iter().product(); + let outer_size = outer_size.max(1); + + let out_shape = reduce_output_shape(shape, dims, keepdim); + let out = Tensor::::empty(&out_shape, dtype, &client.device); + + let a_ptr = a.storage().ptr(); + let out_ptr = out.storage().ptr(); + + dispatch_dtype!(dtype, T => { + unsafe { + kernels::reduce_kernel_with_precision::( + op, + a_ptr as *const T, + out_ptr as *mut T, + reduce_size, + outer_size, + precision, + ); + } + }, op_name); + + Ok(out) + } else if dims.is_empty() { + Ok(a.clone()) + } else if should_fuse_multi_dim_reduction(a, dims) { + reduce_multi_dim_fused(client, op, a, dims, keepdim, precision, op_name) + } else { + let a_contig = ensure_contiguous(a); + + let mut sorted_dims: Vec = dims.to_vec(); + sorted_dims.sort_unstable(); + sorted_dims.reverse(); + + let mut current = a_contig; + for &dim in &sorted_dims { + current = reduce_single_dim_with_precision( + client, op, ¤t, dim, keepdim, precision, op_name, + )?; + } + + Ok(current) + } +} + +/// Reduce a single dimension with explicit accumulation precision. +fn reduce_single_dim_with_precision( + client: &CpuClient, + op: ReduceOp, + a: &Tensor, + dim: usize, + keepdim: bool, + precision: AccumulationPrecision, + op_name: &'static str, +) -> Result> { + let dtype = a.dtype(); + let shape = a.shape(); + let ndim = shape.len(); + + if dim >= ndim { + return Err(Error::InvalidDimension { + dim: dim as isize, + ndim, + }); + } + + let reduce_size = shape[dim]; + let outer_size: usize = shape[..dim].iter().product(); + let outer_size = outer_size.max(1); + let inner_size: usize = shape[dim + 1..].iter().product(); + let inner_size = inner_size.max(1); + + let out_shape = reduce_output_shape(shape, &[dim], keepdim); + let out = Tensor::::empty(&out_shape, dtype, &client.device); + + let a_ptr = a.storage().ptr(); + let out_ptr = out.storage().ptr(); + + if dim == ndim - 1 { + dispatch_dtype!(dtype, T => { + unsafe { + kernels::reduce_kernel_with_precision::( + op, + a_ptr as *const T, + out_ptr as *mut T, + reduce_size, + outer_size, + precision, + ); + } + }, op_name); + } else { + dispatch_dtype!(dtype, T => { + unsafe { + reduce_non_last_dim_with_precision::( + client, + op, + a_ptr as *const T, + out_ptr as *mut T, + outer_size, + reduce_size, + inner_size, + precision, + ); + } + }, op_name); + } + + Ok(out) +} + +#[allow(unsafe_op_in_unsafe_fn)] +unsafe fn reduce_non_last_dim_with_precision( + client: &CpuClient, + op: ReduceOp, + a: *const T, + out: *mut T, + outer_size: usize, + reduce_size: usize, + inner_size: usize, + precision: AccumulationPrecision, +) { + match precision { + AccumulationPrecision::Native => { + // Delegate to the native-precision non-last-dim reduction. + // We inline the serial path here since we already dispatched precision. + for outer in 0..outer_size { + reduce_non_last_dim_outer(op, a, out, outer, reduce_size, inner_size); + } + } + AccumulationPrecision::FP32 | AccumulationPrecision::BF16 => { + reduce_non_last_dim_acc_runtime::( + client, + op, + a, + out, + outer_size, + reduce_size, + inner_size, + ); + } + AccumulationPrecision::FP64 => { + reduce_non_last_dim_acc_runtime::( + client, + op, + a, + out, + outer_size, + reduce_size, + inner_size, + ); + } + } +} + +#[allow(unsafe_op_in_unsafe_fn)] +unsafe fn reduce_non_last_dim_acc( + op: ReduceOp, + a: *const T, + out: *mut T, + outer_size: usize, + reduce_size: usize, + inner_size: usize, +) { + for outer in 0..outer_size { + reduce_non_last_dim_acc_outer::(op, a, out, outer, reduce_size, inner_size); + } +} + +#[allow(unsafe_op_in_unsafe_fn)] +#[inline] +unsafe fn reduce_non_last_dim_acc_outer( + op: ReduceOp, + a: *const T, + out: *mut T, + outer: usize, + reduce_size: usize, + inner_size: usize, +) { + for inner in 0..inner_size { + let first_idx = outer * reduce_size * inner_size + inner; + let first_val = A::acc_in((*a.add(first_idx)).to_f64()); + + let mut acc: A = match op { + ReduceOp::Sum | ReduceOp::Mean => A::ZERO, + ReduceOp::Prod => A::ONE, + ReduceOp::Max | ReduceOp::Min => first_val, + ReduceOp::All => A::ONE, + ReduceOp::Any => A::ZERO, + }; + + for r in 0..reduce_size { + let idx = outer * reduce_size * inner_size + r * inner_size + inner; + let val = A::acc_in((*a.add(idx)).to_f64()); + + acc = match op { + ReduceOp::Sum | ReduceOp::Mean => acc.acc_add(val), + ReduceOp::Prod => acc.acc_mul(val), + ReduceOp::Max => { + if val > acc { + val + } else { + acc + } + } + ReduceOp::Min => { + if val < acc { + val + } else { + acc + } + } + ReduceOp::All => { + if val != A::ZERO && acc != A::ZERO { + A::ONE + } else { + A::ZERO + } + } + ReduceOp::Any => { + if val != A::ZERO || acc != A::ZERO { + A::ONE + } else { + A::ZERO + } + } + }; + } + + if matches!(op, ReduceOp::Mean) { + acc = acc.acc_div(reduce_size); + } + + let out_idx = outer * inner_size + inner; + *out.add(out_idx) = T::from_f64(acc.into()); + } +} + +#[allow(unsafe_op_in_unsafe_fn)] +#[inline] +unsafe fn reduce_non_last_dim_acc_runtime( + client: &CpuClient, + op: ReduceOp, + a: *const T, + out: *mut T, + outer_size: usize, + reduce_size: usize, + inner_size: usize, +) { + #[cfg(feature = "rayon")] + { + if outer_size > 1 { + return reduce_non_last_dim_acc_parallel::( + client, + op, + a, + out, + outer_size, + reduce_size, + inner_size, + ); + } + } + + #[cfg(not(feature = "rayon"))] + let _ = client; + + reduce_non_last_dim_acc::(op, a, out, outer_size, reduce_size, inner_size) +} + +#[cfg(feature = "rayon")] +#[allow(unsafe_op_in_unsafe_fn)] +unsafe fn reduce_non_last_dim_acc_parallel( + client: &CpuClient, + op: ReduceOp, + a: *const T, + out: *mut T, + outer_size: usize, + reduce_size: usize, + inner_size: usize, +) { + let min_len = client.rayon_min_len(); + let a_addr = a as usize; + let out_addr = out as usize; + client.install_parallelism(|| { + (0..outer_size) + .into_par_iter() + .with_min_len(min_len) + .for_each(|outer| unsafe { + let a_ptr = a_addr as *const T; + let out_ptr = out_addr as *mut T; + reduce_non_last_dim_acc_outer::( + op, + a_ptr, + out_ptr, + outer, + reduce_size, + inner_size, + ); + }); + }); +} diff --git a/src/runtime/cpu/helpers/reduce/single_dim.rs b/src/runtime/cpu/helpers/reduce/single_dim.rs new file mode 100644 index 00000000..f2709fa2 --- /dev/null +++ b/src/runtime/cpu/helpers/reduce/single_dim.rs @@ -0,0 +1,223 @@ +//! Single-dimension reduction with native precision + +use crate::dispatch_dtype; +use crate::dtype::Element; +use crate::error::{Error, Result}; +use crate::ops::{Kernel, ReduceOp, reduce_output_shape}; +use crate::runtime::cpu::{CpuClient, CpuRuntime}; +use crate::tensor::Tensor; +#[cfg(feature = "rayon")] +use rayon::prelude::*; + +/// Reduce a single dimension of a tensor using native precision. +/// +/// Uses chunked iteration for non-last dimensions to handle strided memory access. +pub(super) fn reduce_single_dim( + client: &CpuClient, + op: ReduceOp, + a: &Tensor, + dim: usize, + keepdim: bool, + op_name: &'static str, +) -> Result> { + let dtype = a.dtype(); + let shape = a.shape(); + let ndim = shape.len(); + + if dim >= ndim { + return Err(Error::InvalidDimension { + dim: dim as isize, + ndim, + }); + } + + let reduce_size = shape[dim]; + let outer_size: usize = shape[..dim].iter().product(); + let outer_size = outer_size.max(1); + let inner_size: usize = shape[dim + 1..].iter().product(); + let inner_size = inner_size.max(1); + + let out_shape = reduce_output_shape(shape, &[dim], keepdim); + let out = Tensor::::empty(&out_shape, dtype, &client.device); + + if dim == ndim - 1 { + let a_ptr = a.storage().ptr(); + let out_ptr = out.storage().ptr(); + + dispatch_dtype!(dtype, T => { + unsafe { + >::reduce::( + client, + op, + a_ptr as *const T, + out_ptr as *mut T, + reduce_size, + outer_size, + ); + } + }, op_name); + } else { + let a_ptr = a.storage().ptr(); + let out_ptr = out.storage().ptr(); + + dispatch_dtype!(dtype, T => { + unsafe { + reduce_non_last_dim_runtime::( + client, + op, + a_ptr as *const T, + out_ptr as *mut T, + outer_size, + reduce_size, + inner_size, + ); + } + }, op_name); + } + + Ok(out) +} + +#[allow(unsafe_op_in_unsafe_fn)] +unsafe fn reduce_non_last_dim( + op: ReduceOp, + a: *const T, + out: *mut T, + outer_size: usize, + reduce_size: usize, + inner_size: usize, +) { + for outer in 0..outer_size { + reduce_non_last_dim_outer(op, a, out, outer, reduce_size, inner_size); + } +} + +#[allow(unsafe_op_in_unsafe_fn)] +#[inline] +pub(super) unsafe fn reduce_non_last_dim_outer( + op: ReduceOp, + a: *const T, + out: *mut T, + outer: usize, + reduce_size: usize, + inner_size: usize, +) { + for inner in 0..inner_size { + let mut acc = match op { + ReduceOp::Sum | ReduceOp::Mean => T::zero(), + ReduceOp::Prod => T::one(), + ReduceOp::Max => { + let idx = outer * reduce_size * inner_size + inner; + *a.add(idx) + } + ReduceOp::Min => { + let idx = outer * reduce_size * inner_size + inner; + *a.add(idx) + } + ReduceOp::All => T::one(), + ReduceOp::Any => T::zero(), + }; + + for r in 0..reduce_size { + let idx = outer * reduce_size * inner_size + r * inner_size + inner; + let val = *a.add(idx); + + acc = match op { + ReduceOp::Sum | ReduceOp::Mean => acc + val, + ReduceOp::Prod => acc * val, + ReduceOp::Max => { + if val > acc { + val + } else { + acc + } + } + ReduceOp::Min => { + if val < acc { + val + } else { + acc + } + } + ReduceOp::All => { + if val.to_f64() != 0.0 && acc.to_f64() != 0.0 { + T::one() + } else { + T::zero() + } + } + ReduceOp::Any => { + if val.to_f64() != 0.0 || acc.to_f64() != 0.0 { + T::one() + } else { + T::zero() + } + } + }; + } + + if matches!(op, ReduceOp::Mean) { + acc = T::from_f64(acc.to_f64() / reduce_size as f64); + } + + let out_idx = outer * inner_size + inner; + *out.add(out_idx) = acc; + } +} + +#[allow(unsafe_op_in_unsafe_fn)] +unsafe fn reduce_non_last_dim_runtime( + client: &CpuClient, + op: ReduceOp, + a: *const T, + out: *mut T, + outer_size: usize, + reduce_size: usize, + inner_size: usize, +) { + #[cfg(feature = "rayon")] + { + if outer_size > 1 { + return reduce_non_last_dim_parallel( + client, + op, + a, + out, + outer_size, + reduce_size, + inner_size, + ); + } + } + + #[cfg(not(feature = "rayon"))] + let _ = client; + + reduce_non_last_dim(op, a, out, outer_size, reduce_size, inner_size); +} + +#[cfg(feature = "rayon")] +#[allow(unsafe_op_in_unsafe_fn)] +unsafe fn reduce_non_last_dim_parallel( + client: &CpuClient, + op: ReduceOp, + a: *const T, + out: *mut T, + outer_size: usize, + reduce_size: usize, + inner_size: usize, +) { + let min_len = client.rayon_min_len(); + let a_addr = a as usize; + let out_addr = out as usize; + client.install_parallelism(|| { + (0..outer_size) + .into_par_iter() + .with_min_len(min_len) + .for_each(|outer| unsafe { + let a_ptr = a_addr as *const T; + let out_ptr = out_addr as *mut T; + reduce_non_last_dim_outer(op, a_ptr, out_ptr, outer, reduce_size, inner_size); + }); + }); +} diff --git a/src/runtime/cpu/kernels/fft.rs b/src/runtime/cpu/kernels/fft.rs index 10849f05..f2dc5090 100644 --- a/src/runtime/cpu/kernels/fft.rs +++ b/src/runtime/cpu/kernels/fft.rs @@ -129,19 +129,18 @@ pub unsafe fn stockham_fft_batched_c64( batch_size: usize, inverse: bool, normalize_factor: f32, + min_batch_len: usize, ) { use rayon::prelude::*; debug_assert_eq!(input.len(), batch_size * n); debug_assert_eq!(output.len(), batch_size * n); - let output_chunks = std::slice::from_raw_parts_mut(output.as_mut_ptr(), batch_size * n); - // Process batches in parallel - output_chunks - .chunks_mut(n) + output + .par_chunks_mut(n) .enumerate() - .par_bridge() + .with_min_len(min_batch_len.max(1)) .for_each(|(batch_idx, out_chunk)| { let in_start = batch_idx * n; let in_chunk = &input[in_start..in_start + n]; @@ -157,6 +156,7 @@ pub unsafe fn stockham_fft_batched_c64( batch_size: usize, inverse: bool, normalize_factor: f32, + _min_batch_len: usize, ) { for batch_idx in 0..batch_size { let start = batch_idx * n; @@ -256,18 +256,17 @@ pub unsafe fn stockham_fft_batched_c128( batch_size: usize, inverse: bool, normalize_factor: f64, + min_batch_len: usize, ) { use rayon::prelude::*; debug_assert_eq!(input.len(), batch_size * n); debug_assert_eq!(output.len(), batch_size * n); - let output_chunks = std::slice::from_raw_parts_mut(output.as_mut_ptr(), batch_size * n); - - output_chunks - .chunks_mut(n) + output + .par_chunks_mut(n) .enumerate() - .par_bridge() + .with_min_len(min_batch_len.max(1)) .for_each(|(batch_idx, out_chunk)| { let in_start = batch_idx * n; let in_chunk = &input[in_start..in_start + n]; @@ -283,6 +282,7 @@ pub unsafe fn stockham_fft_batched_c128( batch_size: usize, inverse: bool, normalize_factor: f64, + _min_batch_len: usize, ) { for batch_idx in 0..batch_size { let start = batch_idx * n; diff --git a/src/runtime/cpu/mod.rs b/src/runtime/cpu/mod.rs index a31ac1e8..7cf9b118 100644 --- a/src/runtime/cpu/mod.rs +++ b/src/runtime/cpu/mod.rs @@ -35,6 +35,6 @@ pub mod special; pub(crate) mod statistics; pub use crate::tensor::Tensor; -pub use client::{CpuAllocator, CpuClient}; +pub use client::{CpuAllocator, CpuClient, ParallelismConfig}; pub use device::CpuDevice; pub use runtime::CpuRuntime; From 1e0d396772df5184cf6abca7525695e8b1d0b695 Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Tue, 10 Feb 2026 14:56:41 +0800 Subject: [PATCH 4/5] feat(ops): add composite shape operations via impl_generic Implement unfold and repeat_interleave as composite operations built from primitive shape ops (narrow, stack, cat, permute). All backends delegate to impl_generic implementations to guarantee numerical parity and consistent behavior across CPU, CUDA, and WebGPU. This follows the architecture pattern where composite operations share a single algorithm while primitives have backend-specific kernels. --- src/ops/cpu/shape.rs | 20 +++++ src/ops/cuda/shape.rs | 20 +++++ src/ops/impl_generic/mod.rs | 2 + src/ops/impl_generic/shape.rs | 154 ++++++++++++++++++++++++++++++++++ src/ops/traits/shape.rs | 72 ++++++++++++++++ src/ops/wgpu/shape.rs | 20 +++++ 6 files changed, 288 insertions(+) create mode 100644 src/ops/impl_generic/shape.rs diff --git a/src/ops/cpu/shape.rs b/src/ops/cpu/shape.rs index 5886e774..c60d9d7f 100644 --- a/src/ops/cpu/shape.rs +++ b/src/ops/cpu/shape.rs @@ -2,6 +2,7 @@ use crate::error::Result; use crate::ops::ShapeOps; +use crate::ops::impl_generic::{repeat_interleave_impl, unfold_impl}; use crate::runtime::cpu::{ CpuClient, CpuRuntime, helpers::{cat_impl, chunk_impl, pad_impl, repeat_impl, roll_impl, split_impl, stack_impl}, @@ -57,4 +58,23 @@ impl ShapeOps for CpuClient { ) -> Result> { roll_impl(self, tensor, shift, dim) } + + fn unfold( + &self, + tensor: &Tensor, + dim: isize, + size: usize, + step: usize, + ) -> Result> { + unfold_impl(self, tensor, dim, size, step) + } + + fn repeat_interleave( + &self, + tensor: &Tensor, + repeats: usize, + dim: Option, + ) -> Result> { + repeat_interleave_impl(self, tensor, repeats, dim) + } } diff --git a/src/ops/cuda/shape.rs b/src/ops/cuda/shape.rs index d811de41..674b7bcc 100644 --- a/src/ops/cuda/shape.rs +++ b/src/ops/cuda/shape.rs @@ -1,6 +1,7 @@ //! Shape operations for CUDA runtime use crate::error::Result; use crate::ops::ShapeOps; +use crate::ops::impl_generic::{repeat_interleave_impl, unfold_impl}; use crate::runtime::cuda::kernels::{launch_cat_copy, launch_pad, launch_repeat, launch_roll}; use crate::runtime::cuda::{CudaClient, CudaRuntime}; use crate::runtime::{ensure_contiguous, shape_ops}; @@ -182,4 +183,23 @@ impl ShapeOps for CudaClient { Ok(out) } + + fn unfold( + &self, + tensor: &Tensor, + dim: isize, + size: usize, + step: usize, + ) -> Result> { + unfold_impl(self, tensor, dim, size, step) + } + + fn repeat_interleave( + &self, + tensor: &Tensor, + repeats: usize, + dim: Option, + ) -> Result> { + repeat_interleave_impl(self, tensor, repeats, dim) + } } diff --git a/src/ops/impl_generic/mod.rs b/src/ops/impl_generic/mod.rs index ef9a3874..2b6bf006 100644 --- a/src/ops/impl_generic/mod.rs +++ b/src/ops/impl_generic/mod.rs @@ -23,6 +23,7 @@ pub mod einsum; pub mod linalg; pub mod multivariate; pub mod random; +pub mod shape; pub mod utility; #[cfg(any(feature = "cuda", feature = "wgpu"))] @@ -33,6 +34,7 @@ pub use multivariate::{ }; #[cfg(any(feature = "cuda", feature = "wgpu"))] pub use random::randperm_impl; +pub use shape::{repeat_interleave_impl, unfold_impl}; pub use utility::meshgrid_impl; #[cfg(any(feature = "cuda", feature = "wgpu"))] pub use utility::one_hot_impl; diff --git a/src/ops/impl_generic/shape.rs b/src/ops/impl_generic/shape.rs new file mode 100644 index 00000000..02b7f4b4 --- /dev/null +++ b/src/ops/impl_generic/shape.rs @@ -0,0 +1,154 @@ +//! Generic implementations of shape operations (unfold, repeat_interleave). +//! +//! These are composite operations built from primitive shape operations (narrow, stack, cat). +//! All backends delegate to these implementations for numerical parity. + +use crate::error::{Error, Result}; +use crate::ops::ShapeOps; +use crate::runtime::Runtime; +use crate::tensor::Tensor; + +/// Generic unfold implementation. +/// +/// Extracts sliding windows along a dimension. Composed from narrow + stack + permute. +pub fn unfold_impl>( + client: &C, + tensor: &Tensor, + dim: isize, + size: usize, + step: usize, +) -> Result> { + let ndim = tensor.ndim(); + if ndim == 0 { + return Err(Error::InvalidArgument { + arg: "tensor", + reason: "cannot unfold a scalar tensor".to_string(), + }); + } + if size == 0 { + return Err(Error::InvalidArgument { + arg: "size", + reason: "size must be greater than zero".to_string(), + }); + } + if step == 0 { + return Err(Error::InvalidArgument { + arg: "step", + reason: "step must be greater than zero".to_string(), + }); + } + + let dim_idx = if dim < 0 { + let adjusted = ndim as isize + dim; + if adjusted < 0 { + return Err(Error::InvalidDimension { dim, ndim }); + } + adjusted as usize + } else { + dim as usize + }; + if dim_idx >= ndim { + return Err(Error::InvalidDimension { dim, ndim }); + } + + let dim_size = tensor.shape()[dim_idx]; + if size > dim_size { + return Err(Error::InvalidArgument { + arg: "size", + reason: format!( + "size ({}) must be <= dimension {} size ({})", + size, dim_idx, dim_size + ), + }); + } + + let num_windows = (dim_size - size) / step + 1; + + let mut windows = Vec::with_capacity(num_windows); + for i in 0..num_windows { + let start = i * step; + windows.push(tensor.narrow(dim_idx as isize, start, size)?); + } + + let refs: Vec<&Tensor> = windows.iter().collect(); + let stacked = client.stack(&refs, dim_idx as isize)?; + + // stack inserts the window-count axis at dim_idx, and the window-size axis + // lands at dim_idx + 1. Move the size axis to the end to match unfold semantics. + let size_axis = dim_idx + 1; + let out_ndim = stacked.ndim(); + if size_axis + 1 == out_ndim { + return Ok(stacked); + } + + let mut perm = Vec::with_capacity(out_ndim); + for axis in 0..out_ndim { + if axis != size_axis { + perm.push(axis); + } + } + perm.push(size_axis); + stacked.permute(&perm) +} + +/// Generic repeat_interleave implementation. +/// +/// Repeats each element along a dimension. Composed from narrow + cat. +pub fn repeat_interleave_impl>( + client: &C, + tensor: &Tensor, + repeats: usize, + dim: Option, +) -> Result> { + if repeats == 0 { + return Err(Error::InvalidArgument { + arg: "repeats", + reason: "repeats must be greater than zero".to_string(), + }); + } + + let (input, dim_idx) = match dim { + Some(d) => { + let ndim = tensor.ndim(); + if ndim == 0 { + return Err(Error::InvalidDimension { dim: d, ndim: 0 }); + } + let dim_idx = if d < 0 { + let adjusted = ndim as isize + d; + if adjusted < 0 { + return Err(Error::InvalidDimension { dim: d, ndim }); + } + adjusted as usize + } else { + d as usize + }; + if dim_idx >= ndim { + return Err(Error::InvalidDimension { dim: d, ndim }); + } + (tensor.clone(), dim_idx) + } + None => (tensor.contiguous().flatten()?, 0usize), + }; + + let dim_size = input.shape()[dim_idx]; + if dim_size == 0 { + let mut out_shape = input.shape().to_vec(); + out_shape[dim_idx] = 0; + return Ok(Tensor::::empty( + &out_shape, + input.dtype(), + input.device(), + )); + } + + let mut chunks: Vec> = Vec::with_capacity(dim_size * repeats); + for i in 0..dim_size { + let slice = input.narrow(dim_idx as isize, i, 1)?; + for _ in 0..repeats { + chunks.push(slice.clone()); + } + } + + let refs: Vec<&Tensor> = chunks.iter().collect(); + client.cat(&refs, dim_idx as isize) +} diff --git a/src/ops/traits/shape.rs b/src/ops/traits/shape.rs index 0ff7f30a..d9026dba 100644 --- a/src/ops/traits/shape.rs +++ b/src/ops/traits/shape.rs @@ -211,4 +211,76 @@ pub trait ShapeOps { /// # Ok::<(), numr::error::Error>(()) /// ``` fn roll(&self, tensor: &Tensor, shift: isize, dim: isize) -> Result>; + + /// Extract sliding local windows along a dimension. + /// + /// Returns a tensor containing all windows of length `size` sampled every `step` + /// elements along `dim`. The output has one extra dimension, with the window-size + /// dimension appended at the end. + /// + /// If input shape is `` `[d0, ..., d_dim, ..., dn]` ``, output shape is + /// `` `[d0, ..., num_windows, ..., dn, size]` `` where: + /// `` `num_windows = (d_dim - size) / step + 1` ``. + /// + /// # Arguments + /// + /// * `tensor` - Input tensor + /// * `dim` - Dimension along which to extract windows (supports negative indexing) + /// * `size` - Window size (must be > 0 and <= dimension size) + /// * `step` - Stride between window starts (must be > 0) + /// + /// # Returns + /// + /// New tensor containing extracted windows + /// + /// # Example + /// + /// ``` + /// # use numr::prelude::*; + /// # let device = CpuDevice::new(); + /// # let client = CpuRuntime::default_client(&device); + /// use numr::ops::ShapeOps; + /// + /// let a = Tensor::::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0], &[5], &device); + /// let windows = client.unfold(&a, 0, 3, 1)?; // Shape: [3, 3] + /// // Result: [[1,2,3], [2,3,4], [3,4,5]] + /// # Ok::<(), numr::error::Error>(()) + /// ``` + fn unfold(&self, tensor: &Tensor, dim: isize, size: usize, step: usize) + -> Result>; + + /// Repeat each element along a dimension. + /// + /// Unlike `repeat`, which tiles whole tensor blocks along each dimension, + /// `repeat_interleave` repeats individual elements in-place along one dimension. + /// + /// # Arguments + /// + /// * `tensor` - Input tensor + /// * `repeats` - Number of times to repeat each element (must be > 0) + /// * `dim` - Dimension to repeat along (supports negative indexing). If `None`, input is flattened first. + /// + /// # Returns + /// + /// New tensor with repeated elements + /// + /// # Example + /// + /// ``` + /// # use numr::prelude::*; + /// # let device = CpuDevice::new(); + /// # let client = CpuRuntime::default_client(&device); + /// use numr::ops::ShapeOps; + /// + /// let a = Tensor::::from_slice(&[1.0f32, 2.0, 3.0], &[3], &device); + /// let out = client.repeat_interleave(&a, 2, Some(0))?; + /// // Result: [1, 1, 2, 2, 3, 3] + /// # Ok::<(), numr::error::Error>(()) + /// ``` + fn repeat_interleave( + &self, + tensor: &Tensor, + repeats: usize, + dim: Option, + ) -> Result>; } diff --git a/src/ops/wgpu/shape.rs b/src/ops/wgpu/shape.rs index c4d70031..dd2ece85 100644 --- a/src/ops/wgpu/shape.rs +++ b/src/ops/wgpu/shape.rs @@ -3,6 +3,7 @@ use crate::dtype::DType; use crate::error::{Error, Result}; use crate::ops::ShapeOps; +use crate::ops::impl_generic::{repeat_interleave_impl, unfold_impl}; use crate::runtime::shape_ops; use crate::runtime::shape_ops::{validate_cat, validate_stack}; use crate::runtime::wgpu::WgpuClient; @@ -361,4 +362,23 @@ impl ShapeOps for WgpuClient { Ok(out) } + + fn unfold( + &self, + tensor: &Tensor, + dim: isize, + size: usize, + step: usize, + ) -> Result> { + unfold_impl(self, tensor, dim, size, step) + } + + fn repeat_interleave( + &self, + tensor: &Tensor, + repeats: usize, + dim: Option, + ) -> Result> { + repeat_interleave_impl(self, tensor, repeats, dim) + } } From a6405e26e40c2e4a14a5674bcb0bc7ef9537eabd Mon Sep 17 00:00:00 2001 From: Farhan Syah Date: Tue, 10 Feb 2026 14:57:07 +0800 Subject: [PATCH 5/5] refactor: add configurable parallelism and improve kernel consistency Introduce configurable chunk_size parameter across CPU kernels to allow tuning parallelism granularity. This enables better performance on diverse workloads and hardware configurations. Changes: - Add chunk_size parameter to complex, FFT, and other parallel kernels - Feature-gate half_convert module for x86_64 only - Update indexing operations for consistency across backends - Add ParallelismConfig tests to verify chunking behavior - Update backend parity tests for new configurations No functional changes to existing default behavior. --- src/ops/cpu/complex.rs | 219 +++++++++----- src/ops/cpu/matmul.rs | 103 +++++++ src/ops/cuda/indexing/advanced.rs | 137 ++++++--- src/ops/cuda/indexing/gather_scatter.rs | 106 +++---- src/ops/cuda/indexing/helpers.rs | 21 +- src/ops/traits/indexing.rs | 92 +++++- src/ops/wgpu/indexing.rs | 168 ++++++++--- src/runtime/cpu/helpers/indexing.rs | 179 +++++------- src/runtime/cpu/kernels/complex.rs | 185 +++++++----- src/runtime/cpu/kernels/simd/matmul/mod.rs | 1 + src/runtime/cuda/kernels/index.cu | 135 +++++++++ src/runtime/cuda/kernels/index.rs | 147 ++++++++++ src/runtime/wgpu/ops/helpers.rs | 10 + src/runtime/wgpu/shaders/generator/index.rs | 228 +++++++++++++++ src/runtime/wgpu/shaders/generator/mod.rs | 5 +- src/runtime/wgpu/shaders/index.rs | 159 +++++++++- src/runtime/wgpu/shaders/mod.rs | 5 +- tests/backend_parity/complex.rs | 96 ++++++- tests/backend_parity/fft.rs | 119 +++++++- tests/backend_parity/indexing.rs | 172 +++++++++++ tests/backend_parity/indexing_advanced.rs | 303 ++++++++++++++++++++ tests/backend_parity/matmul.rs | 29 ++ tests/backend_parity/matmul_bias.rs | 35 +++ tests/backend_parity/reduce.rs | 86 ++++++ tests/backend_parity/shape.rs | 287 ++++++++++++++++++ 25 files changed, 2608 insertions(+), 419 deletions(-) diff --git a/src/ops/cpu/complex.rs b/src/ops/cpu/complex.rs index 41c0f1bb..06d4b9bf 100644 --- a/src/ops/cpu/complex.rs +++ b/src/ops/cpu/complex.rs @@ -29,17 +29,30 @@ impl ComplexOps for CpuClient { let a_ptr = a_contig.storage().ptr(); let out_ptr = out.storage().ptr(); + let chunk_size = self.chunk_size_hint(); - unsafe { - match dtype { - DType::Complex64 => { - kernels::conj_complex64(a_ptr as *const _, out_ptr as *mut _, numel); - } - DType::Complex128 => { - kernels::conj_complex128(a_ptr as *const _, out_ptr as *mut _, numel); - } - _ => unreachable!("conj called on non-complex dtype"), + match dtype { + DType::Complex64 => { + self.install_parallelism(|| unsafe { + kernels::conj_complex64( + a_ptr as *const _, + out_ptr as *mut _, + numel, + chunk_size, + ); + }); + } + DType::Complex128 => { + self.install_parallelism(|| unsafe { + kernels::conj_complex128( + a_ptr as *const _, + out_ptr as *mut _, + numel, + chunk_size, + ); + }); } + _ => unreachable!("conj called on non-complex dtype"), } Ok(out) @@ -71,17 +84,30 @@ impl ComplexOps for CpuClient { let a_ptr = a_contig.storage().ptr(); let out_ptr = out.storage().ptr(); + let chunk_size = self.chunk_size_hint(); - unsafe { - match dtype { - DType::Complex64 => { - kernels::real_complex64(a_ptr as *const _, out_ptr as *mut _, numel); - } - DType::Complex128 => { - kernels::real_complex128(a_ptr as *const _, out_ptr as *mut _, numel); - } - _ => unreachable!("real called on non-complex dtype"), + match dtype { + DType::Complex64 => { + self.install_parallelism(|| unsafe { + kernels::real_complex64( + a_ptr as *const _, + out_ptr as *mut _, + numel, + chunk_size, + ); + }); } + DType::Complex128 => { + self.install_parallelism(|| unsafe { + kernels::real_complex128( + a_ptr as *const _, + out_ptr as *mut _, + numel, + chunk_size, + ); + }); + } + _ => unreachable!("real called on non-complex dtype"), } Ok(out) @@ -113,17 +139,30 @@ impl ComplexOps for CpuClient { let a_ptr = a_contig.storage().ptr(); let out_ptr = out.storage().ptr(); + let chunk_size = self.chunk_size_hint(); - unsafe { - match dtype { - DType::Complex64 => { - kernels::imag_complex64(a_ptr as *const _, out_ptr as *mut _, numel); - } - DType::Complex128 => { - kernels::imag_complex128(a_ptr as *const _, out_ptr as *mut _, numel); - } - _ => unreachable!("imag called on non-complex dtype"), + match dtype { + DType::Complex64 => { + self.install_parallelism(|| unsafe { + kernels::imag_complex64( + a_ptr as *const _, + out_ptr as *mut _, + numel, + chunk_size, + ); + }); } + DType::Complex128 => { + self.install_parallelism(|| unsafe { + kernels::imag_complex128( + a_ptr as *const _, + out_ptr as *mut _, + numel, + chunk_size, + ); + }); + } + _ => unreachable!("imag called on non-complex dtype"), } Ok(out) @@ -133,6 +172,7 @@ impl ComplexOps for CpuClient { let dtype = a.dtype(); let shape = a.shape(); let numel = a.numel(); + let chunk_size = self.chunk_size_hint(); let a_contig = ensure_contiguous(a); @@ -148,19 +188,31 @@ impl ComplexOps for CpuClient { let a_ptr = a_contig.storage().ptr(); let out_ptr = out.storage().ptr(); - unsafe { - match dtype { - DType::F32 => { - kernels::angle_real_f32(a_ptr as *const _, out_ptr as *mut _, numel); - } - DType::F64 => { - kernels::angle_real_f64(a_ptr as *const _, out_ptr as *mut _, numel); - } - _ => { - // For integer types, angle doesn't make mathematical sense - // Return zeros - return Ok(Tensor::::zeros(shape, dtype, &self.device)); - } + match dtype { + DType::F32 => { + self.install_parallelism(|| unsafe { + kernels::angle_real_f32( + a_ptr as *const _, + out_ptr as *mut _, + numel, + chunk_size, + ); + }); + } + DType::F64 => { + self.install_parallelism(|| unsafe { + kernels::angle_real_f64( + a_ptr as *const _, + out_ptr as *mut _, + numel, + chunk_size, + ); + }); + } + _ => { + // For integer types, angle doesn't make mathematical sense + // Return zeros + return Ok(Tensor::::zeros(shape, dtype, &self.device)); } } return Ok(out); @@ -181,16 +233,28 @@ impl ComplexOps for CpuClient { let a_ptr = a_contig.storage().ptr(); let out_ptr = out.storage().ptr(); - unsafe { - match dtype { - DType::Complex64 => { - kernels::angle_complex64(a_ptr as *const _, out_ptr as *mut _, numel); - } - DType::Complex128 => { - kernels::angle_complex128(a_ptr as *const _, out_ptr as *mut _, numel); - } - _ => unreachable!("angle called on non-complex dtype"), + match dtype { + DType::Complex64 => { + self.install_parallelism(|| unsafe { + kernels::angle_complex64( + a_ptr as *const _, + out_ptr as *mut _, + numel, + chunk_size, + ); + }); } + DType::Complex128 => { + self.install_parallelism(|| unsafe { + kernels::angle_complex128( + a_ptr as *const _, + out_ptr as *mut _, + numel, + chunk_size, + ); + }); + } + _ => unreachable!("angle called on non-complex dtype"), } Ok(out) @@ -226,27 +290,32 @@ impl ComplexOps for CpuClient { let real_ptr = real_contig.storage().ptr(); let imag_ptr = imag_contig.storage().ptr(); let out_ptr = out.storage().ptr(); + let chunk_size = self.chunk_size_hint(); - unsafe { - match input_dtype { - DType::F32 => { + match input_dtype { + DType::F32 => { + self.install_parallelism(|| unsafe { kernels::from_real_imag_f32( real_ptr as *const _, imag_ptr as *const _, out_ptr as *mut _, numel, + chunk_size, ); - } - DType::F64 => { + }); + } + DType::F64 => { + self.install_parallelism(|| unsafe { kernels::from_real_imag_f64( real_ptr as *const _, imag_ptr as *const _, out_ptr as *mut _, numel, + chunk_size, ); - } - _ => unreachable!("validated above"), + }); } + _ => unreachable!("validated above"), } Ok(out) @@ -275,27 +344,32 @@ impl ComplexOps for CpuClient { let complex_ptr = complex_contig.storage().ptr(); let real_ptr = real_contig.storage().ptr(); let out_ptr = out.storage().ptr(); + let chunk_size = self.chunk_size_hint(); - unsafe { - match dtype { - DType::Complex64 => { + match dtype { + DType::Complex64 => { + self.install_parallelism(|| unsafe { kernels::complex64_mul_real( complex_ptr as *const _, real_ptr as *const _, out_ptr as *mut _, numel, + chunk_size, ); - } - DType::Complex128 => { + }); + } + DType::Complex128 => { + self.install_parallelism(|| unsafe { kernels::complex128_mul_real( complex_ptr as *const _, real_ptr as *const _, out_ptr as *mut _, numel, + chunk_size, ); - } - _ => unreachable!("validated above"), + }); } + _ => unreachable!("validated above"), } Ok(out) @@ -324,27 +398,32 @@ impl ComplexOps for CpuClient { let complex_ptr = complex_contig.storage().ptr(); let real_ptr = real_contig.storage().ptr(); let out_ptr = out.storage().ptr(); + let chunk_size = self.chunk_size_hint(); - unsafe { - match dtype { - DType::Complex64 => { + match dtype { + DType::Complex64 => { + self.install_parallelism(|| unsafe { kernels::complex64_div_real( complex_ptr as *const _, real_ptr as *const _, out_ptr as *mut _, numel, + chunk_size, ); - } - DType::Complex128 => { + }); + } + DType::Complex128 => { + self.install_parallelism(|| unsafe { kernels::complex128_div_real( complex_ptr as *const _, real_ptr as *const _, out_ptr as *mut _, numel, + chunk_size, ); - } - _ => unreachable!("validated above"), + }); } + _ => unreachable!("validated above"), } Ok(out) diff --git a/src/ops/cpu/matmul.rs b/src/ops/cpu/matmul.rs index afac7b9c..8b94693c 100644 --- a/src/ops/cpu/matmul.rs +++ b/src/ops/cpu/matmul.rs @@ -66,6 +66,57 @@ impl MatmulOps for CpuClient { // Dispatch based on dtype dispatch_dtype!(dtype, T => { + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + + if batch_size > 1 { + let min_len = self.rayon_min_len(); + self.install_parallelism(|| { + (0..batch_size) + .into_par_iter() + .with_min_len(min_len) + .for_each(|batch| unsafe { + let a_offset = batch * m * k; + let b_offset = batch * k * n; + let out_offset = batch * m * n; + + >::matmul::( + self, + (a_ptr as *const T).add(a_offset), + (b_ptr as *const T).add(b_offset), + (out_ptr as *mut T).add(out_offset), + m, + n, + k, + lda, + ldb, + ldc, + ); + }); + }); + } else { + unsafe { + let a_offset = 0; + let b_offset = 0; + let out_offset = 0; + >::matmul::( + self, + (a_ptr as *const T).add(a_offset), + (b_ptr as *const T).add(b_offset), + (out_ptr as *mut T).add(out_offset), + m, + n, + k, + lda, + ldb, + ldc, + ); + } + } + } + + #[cfg(not(feature = "rayon"))] unsafe { for batch in 0..batch_size { let a_offset = batch * m * k; @@ -149,6 +200,58 @@ impl MatmulOps for CpuClient { // Dispatch based on dtype dispatch_dtype!(dtype, T => { + #[cfg(feature = "rayon")] + { + use rayon::prelude::*; + + if batch_size > 1 { + let min_len = self.rayon_min_len(); + self.install_parallelism(|| { + (0..batch_size) + .into_par_iter() + .with_min_len(min_len) + .for_each(|batch| unsafe { + let a_offset = batch * m * k; + let b_offset = batch * k * n; + let out_offset = batch * m * n; + + matmul_bias_kernel::( + (a_ptr as *const T).add(a_offset), + (b_ptr as *const T).add(b_offset), + bias_ptr as *const T, // bias is 1D, same for all batches + (out_ptr as *mut T).add(out_offset), + m, + n, + k, + lda, + ldb, + ldc, + ); + }); + }); + } else { + unsafe { + let a_offset = 0; + let b_offset = 0; + let out_offset = 0; + + matmul_bias_kernel::( + (a_ptr as *const T).add(a_offset), + (b_ptr as *const T).add(b_offset), + bias_ptr as *const T, + (out_ptr as *mut T).add(out_offset), + m, + n, + k, + lda, + ldb, + ldc, + ); + } + } + } + + #[cfg(not(feature = "rayon"))] unsafe { for batch in 0..batch_size { let a_offset = batch * m * k; diff --git a/src/ops/cuda/indexing/advanced.rs b/src/ops/cuda/indexing/advanced.rs index 423abd37..72473f83 100644 --- a/src/ops/cuda/indexing/advanced.rs +++ b/src/ops/cuda/indexing/advanced.rs @@ -5,12 +5,15 @@ use crate::error::{Error, Result}; use crate::ops::{ReduceOps, ScatterReduceOp, TypeConversionOps}; use crate::runtime::cuda::kernels::{ ScatterReduceOpCuda, launch_bincount_weighted, launch_copy, launch_embedding_lookup, - launch_fill_with_f64, launch_gather_nd, launch_scatter_reduce, + launch_fill_with_f64, launch_gather_nd, launch_scatter_reduce, launch_scatter_reduce_count, + launch_scatter_reduce_mean_div, }; use crate::runtime::cuda::{CudaClient, CudaRuntime}; use crate::runtime::{Runtime, compute_contiguous_strides, ensure_contiguous}; use crate::tensor::Tensor; +use super::helpers::normalize_indices_to_i64; + /// Execute embedding_lookup operation. pub fn embedding_lookup( client: &CudaClient, @@ -28,24 +31,18 @@ pub fn embedding_lookup( }); } - // Validate indices dtype - if indices.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: indices.dtype(), - }); - } + let indices_i64 = normalize_indices_to_i64(client, indices)?; let vocab_size = emb_shape[0]; let embedding_dim = emb_shape[1]; - let num_indices = indices.numel(); + let num_indices = indices_i64.numel(); // Output shape: indices.shape() + [embedding_dim] - let mut out_shape = indices.shape().to_vec(); + let mut out_shape = indices_i64.shape().to_vec(); out_shape.push(embedding_dim); let emb_contig = ensure_contiguous(embeddings); - let idx_contig = ensure_contiguous(indices); + let idx_contig = ensure_contiguous(&indices_i64); let out = Tensor::::empty(&out_shape, dtype, &client.device); unsafe { @@ -88,13 +85,7 @@ pub fn scatter_reduce( }); } - // Validate dtypes - if index.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: index.dtype(), - }); - } + let index_i64 = normalize_indices_to_i64(client, index)?; if src.dtype() != dtype { return Err(Error::DTypeMismatch { @@ -104,35 +95,32 @@ pub fn scatter_reduce( } // Validate that index and src have same shape - if index.shape() != src.shape() { + if index_i64.shape() != src.shape() { return Err(Error::ShapeMismatch { expected: src.shape().to_vec(), - got: index.shape().to_vec(), + got: index_i64.shape().to_vec(), }); } // Validate that index has same number of dimensions as dst - if index.ndim() != ndim { + if index_i64.ndim() != ndim { return Err(Error::ShapeMismatch { expected: shape.to_vec(), - got: index.shape().to_vec(), + got: index_i64.shape().to_vec(), }); } // Map ScatterReduceOp to ScatterReduceOpCuda let cuda_op = match op { - ScatterReduceOp::Sum | ScatterReduceOp::Mean => ScatterReduceOpCuda::Sum, + ScatterReduceOp::Sum => ScatterReduceOpCuda::Sum, ScatterReduceOp::Max => ScatterReduceOpCuda::Max, ScatterReduceOp::Min => ScatterReduceOpCuda::Min, - ScatterReduceOp::Prod => { - return Err(Error::NotImplemented { - feature: "scatter_reduce with Prod on CUDA", - }); - } + ScatterReduceOp::Prod => ScatterReduceOpCuda::Prod, + ScatterReduceOp::Mean => ScatterReduceOpCuda::Sum, // Mean uses sum kernel + count + div }; let dst_contig = ensure_contiguous(dst); - let index_contig = ensure_contiguous(index); + let index_contig = ensure_contiguous(&index_i64); let src_contig = ensure_contiguous(src); // Allocate output and initialize with dst values if include_self @@ -196,12 +184,78 @@ pub fn scatter_reduce( )?; } - // For Mean operation, we would need to track counts and divide - // This is not currently supported in the kernel - if op == ScatterReduceOp::Mean { - return Err(Error::NotImplemented { - feature: "scatter_reduce with Mean on CUDA (requires count tracking)", - }); + // For mean: divide sum by count + if matches!(op, ScatterReduceOp::Mean) { + // Only float types support mean + if !matches!(dtype, DType::F32 | DType::F64) { + return Err(Error::UnsupportedDType { + dtype, + op: "scatter_reduce_mean", + }); + } + + // Allocate count buffer (same shape as output, zero-initialized) + let count = Tensor::::empty(shape, dtype, &client.device); + unsafe { + launch_fill_with_f64( + &client.context, + &client.stream, + client.device.index, + dtype, + 0.0, + count.storage().ptr(), + dst.numel(), + )?; + } + + // If include_self, each dst element starts with count=1 + if include_self { + unsafe { + launch_fill_with_f64( + &client.context, + &client.stream, + client.device.index, + dtype, + 1.0, + count.storage().ptr(), + dst.numel(), + )?; + } + } + + // Scatter count: atomicAdd 1 for each src element + unsafe { + launch_scatter_reduce_count( + &client.context, + &client.stream, + client.device.index, + dtype, + index_contig.storage().ptr(), + count.storage().ptr(), + dim, + outer_size, + dim_size, + inner_size, + src_dim_size, + )?; + } + + // Divide sum by count + let result = Tensor::::empty(shape, dtype, &client.device); + unsafe { + launch_scatter_reduce_mean_div( + &client.context, + &client.stream, + client.device.index, + dtype, + out.storage().ptr(), + count.storage().ptr(), + result.storage().ptr(), + dst.numel(), + )?; + } + + return Ok(result); } Ok(out) @@ -215,15 +269,8 @@ pub fn gather_nd( ) -> Result> { let dtype = input.dtype(); let input_shape = input.shape(); - let indices_shape = indices.shape(); - - // Validate indices dtype - if indices.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: indices.dtype(), - }); - } + let indices_i64 = normalize_indices_to_i64(client, indices)?; + let indices_shape = indices_i64.shape(); // Indices must have at least 1 dimension if indices_shape.is_empty() { @@ -263,7 +310,7 @@ pub fn gather_nd( let slice_size = slice_size.max(1); let input_contig = ensure_contiguous(input); - let indices_contig = ensure_contiguous(indices); + let indices_contig = ensure_contiguous(&indices_i64); let out = Tensor::::empty(&out_shape, dtype, &client.device); // Allocate device memory for input shape and strides diff --git a/src/ops/cuda/indexing/gather_scatter.rs b/src/ops/cuda/indexing/gather_scatter.rs index 505ef990..c4be89a7 100644 --- a/src/ops/cuda/indexing/gather_scatter.rs +++ b/src/ops/cuda/indexing/gather_scatter.rs @@ -10,6 +10,8 @@ use crate::runtime::cuda::{CudaClient, CudaRuntime}; use crate::runtime::{Runtime, compute_contiguous_strides, ensure_contiguous}; use crate::tensor::Tensor; +use super::helpers::normalize_indices_to_i64; + /// Execute gather operation along dimension. pub fn gather( client: &CudaClient, @@ -17,13 +19,7 @@ pub fn gather( dim: usize, index: &Tensor, ) -> Result> { - // Validate index dtype - if index.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: index.dtype(), - }); - } + let index_i64 = normalize_indices_to_i64(client, index)?; // Validate dimension let ndim = a.ndim(); @@ -35,19 +31,19 @@ pub fn gather( } // Validate index tensor has same number of dimensions - if index.ndim() != ndim { + if index_i64.ndim() != ndim { return Err(Error::ShapeMismatch { expected: a.shape().to_vec(), - got: index.shape().to_vec(), + got: index_i64.shape().to_vec(), }); } let dtype = a.dtype(); let a_contig = ensure_contiguous(a); - let index_contig = ensure_contiguous(index); + let index_contig = ensure_contiguous(&index_i64); // Output has same shape as index - let out_shape = index.shape().to_vec(); + let out_shape = index_i64.shape().to_vec(); let out = Tensor::::empty(&out_shape, dtype, &client.device); // Prepare shape and stride arrays for GPU @@ -117,13 +113,7 @@ pub fn scatter( index: &Tensor, src: &Tensor, ) -> Result> { - // Validate index dtype - if index.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: index.dtype(), - }); - } + let index_i64 = normalize_indices_to_i64(client, index)?; // Validate dimension let ndim = a.ndim(); @@ -144,15 +134,15 @@ pub fn scatter( } // Index and src must have same shape - if index.shape() != src.shape() { + if index_i64.shape() != src.shape() { return Err(Error::ShapeMismatch { - expected: index.shape().to_vec(), + expected: index_i64.shape().to_vec(), got: src.shape().to_vec(), }); } let a_contig = ensure_contiguous(a); - let index_contig = ensure_contiguous(index); + let index_contig = ensure_contiguous(&index_i64); let src_contig = ensure_contiguous(src); // Output has same shape as input @@ -249,19 +239,13 @@ pub fn index_select( dim: usize, index: &Tensor, ) -> Result> { - // Validate index dtype - if index.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: index.dtype(), - }); - } + let index_i64 = normalize_indices_to_i64(client, index)?; // Validate index is 1D - if index.ndim() != 1 { + if index_i64.ndim() != 1 { return Err(Error::ShapeMismatch { - expected: vec![index.numel()], - got: index.shape().to_vec(), + expected: vec![index_i64.numel()], + got: index_i64.shape().to_vec(), }); } @@ -277,10 +261,10 @@ pub fn index_select( let dtype = a.dtype(); let a_contig = ensure_contiguous(a); - let index_contig = ensure_contiguous(index); + let index_contig = ensure_contiguous(&index_i64); // Compute output shape: same as input but dim[dim] = index.len() - let index_len = index.numel(); + let index_len = index_i64.numel(); let mut out_shape = shape.to_vec(); out_shape[dim] = index_len; @@ -373,48 +357,36 @@ pub fn gather_2d( let nrows = shape[0]; let ncols = shape[1]; - // Validate index dtypes - if rows.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: rows.dtype(), - }); - } - - if cols.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: cols.dtype(), - }); - } + let rows_i64 = normalize_indices_to_i64(client, rows)?; + let cols_i64 = normalize_indices_to_i64(client, cols)?; // Validate rows and cols are 1D and have same length - if rows.ndim() != 1 { + if rows_i64.ndim() != 1 { return Err(Error::ShapeMismatch { - expected: vec![rows.numel()], - got: rows.shape().to_vec(), + expected: vec![rows_i64.numel()], + got: rows_i64.shape().to_vec(), }); } - if cols.ndim() != 1 { + if cols_i64.ndim() != 1 { return Err(Error::ShapeMismatch { - expected: vec![cols.numel()], - got: cols.shape().to_vec(), + expected: vec![cols_i64.numel()], + got: cols_i64.shape().to_vec(), }); } - let num_indices = rows.numel(); - if cols.numel() != num_indices { + let num_indices = rows_i64.numel(); + if cols_i64.numel() != num_indices { return Err(Error::ShapeMismatch { expected: vec![num_indices], - got: cols.shape().to_vec(), + got: cols_i64.shape().to_vec(), }); } // Make all inputs contiguous let input_contig = ensure_contiguous(input); - let rows_contig = ensure_contiguous(rows); - let cols_contig = ensure_contiguous(cols); + let rows_contig = ensure_contiguous(&rows_i64); + let cols_contig = ensure_contiguous(&cols_i64); // Allocate output let out = Tensor::::empty(&[num_indices], dtype, &client.device); @@ -458,19 +430,13 @@ pub fn index_put( }); } - // Validate index dtype - if index.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: index.dtype(), - }); - } + let index_i64 = normalize_indices_to_i64(client, index)?; // Validate index is 1D - if index.ndim() != 1 { + if index_i64.ndim() != 1 { return Err(Error::ShapeMismatch { - expected: vec![index.numel()], - got: index.shape().to_vec(), + expected: vec![index_i64.numel()], + got: index_i64.shape().to_vec(), }); } @@ -482,7 +448,7 @@ pub fn index_put( }); } - let index_len = index.numel(); + let index_len = index_i64.numel(); // Validate src shape: must match a's shape except at dim where it equals index_len let mut expected_src_shape = shape.to_vec(); @@ -495,7 +461,7 @@ pub fn index_put( } let a_contig = ensure_contiguous(a); - let index_contig = ensure_contiguous(index); + let index_contig = ensure_contiguous(&index_i64); let src_contig = ensure_contiguous(src); // Compute dim_size for validation diff --git a/src/ops/cuda/indexing/helpers.rs b/src/ops/cuda/indexing/helpers.rs index 9e6d1567..981533c8 100644 --- a/src/ops/cuda/indexing/helpers.rs +++ b/src/ops/cuda/indexing/helpers.rs @@ -2,9 +2,10 @@ use crate::dtype::DType; use crate::error::{Error, Result}; +use crate::ops::TypeConversionOps; use crate::ops::broadcast_shape; use crate::runtime::cuda::kernels::compute_broadcast_strides; -use crate::runtime::cuda::{CudaDevice, CudaRuntime}; +use crate::runtime::cuda::{CudaClient, CudaDevice, CudaRuntime}; use crate::tensor::Tensor; /// Validates that the mask tensor has dtype U8. @@ -18,6 +19,24 @@ pub fn validate_mask_dtype(mask: &Tensor) -> Result<()> { Ok(()) } +/// Normalize index tensor to I64 for CUDA indexing kernels. +/// +/// CUDA indexing kernels consume i64 indices. This helper accepts both I32 and I64 +/// public API inputs and casts to I64 on device when needed. +pub fn normalize_indices_to_i64( + client: &CudaClient, + indices: &Tensor, +) -> Result> { + match indices.dtype() { + DType::I64 => Ok(indices.clone()), + DType::I32 => client.cast(indices, DType::I64), + other => Err(Error::DTypeMismatch { + lhs: DType::I64, + rhs: other, + }), + } +} + /// Context for broadcast masked operations. /// Holds the GPU tensors needed for broadcast stride computation. pub struct BroadcastContext { diff --git a/src/ops/traits/indexing.rs b/src/ops/traits/indexing.rs index 68cd2413..2c3515a6 100644 --- a/src/ops/traits/indexing.rs +++ b/src/ops/traits/indexing.rs @@ -1,5 +1,6 @@ //! Indexing operations trait. +use crate::dtype::DType; use crate::error::{Error, Result}; use crate::runtime::Runtime; use crate::tensor::Tensor; @@ -22,6 +23,20 @@ pub enum ScatterReduceOp { Prod, } +/// Validate that indices tensor has an integer dtype (I32 or I64). +fn validate_index_dtype(indices: &Tensor) -> Result<()> { + match indices.dtype() { + DType::I32 | DType::I64 => Ok(()), + other => Err(Error::InvalidArgument { + arg: "indices", + reason: format!( + "indices must have integer dtype (I32 or I64), got {:?}", + other + ), + }), + } +} + /// Indexing operations pub trait IndexingOps { /// Argmax: returns indices of maximum values along a dimension. @@ -77,7 +92,7 @@ pub trait IndexingOps { /// /// * `a` - Input tensor /// * `dim` - Dimension along which to gather - /// * `index` - Index tensor (I64) with same number of dimensions as input + /// * `index` - Index tensor (I32 or I64) with same number of dimensions as input /// /// # Returns /// @@ -101,7 +116,7 @@ pub trait IndexingOps { /// /// * `a` - Input tensor (values to scatter into) /// * `dim` - Dimension along which to scatter - /// * `index` - Index tensor (I64) specifying scatter positions + /// * `index` - Index tensor (I32 or I64) specifying scatter positions /// * `src` - Source tensor with values to scatter /// /// # Returns @@ -129,7 +144,7 @@ pub trait IndexingOps { /// /// * `a` - Input tensor /// * `dim` - Dimension along which to select - /// * `index` - 1D index tensor (I64) of length m + /// * `index` - 1D index tensor (I32 or I64) of length m /// /// # Returns /// @@ -175,7 +190,7 @@ pub trait IndexingOps { /// /// * `a` - Input tensor to modify (copied, not mutated) /// * `dim` - Dimension along which to put values - /// * `index` - 1D index tensor (I64) specifying positions + /// * `index` - 1D index tensor (I32 or I64) specifying positions /// * `src` - Source tensor with values to insert. Shape must match `a` except /// at `dim` where it must equal `index.numel()` /// @@ -200,6 +215,63 @@ pub trait IndexingOps { }) } + /// Take values from a tensor using flat indices. + /// + /// The input tensor is treated as flattened 1D storage, and values are gathered + /// at positions specified by `indices`. The output shape matches `indices.shape()`. + /// + /// # Arguments + /// + /// * `tensor` - Input tensor to gather from + /// * `indices` - Index tensor (I32 or I64) containing flat indices + /// + /// # Returns + /// + /// Tensor of shape `indices.shape()` with gathered values + fn take(&self, tensor: &Tensor, indices: &Tensor) -> Result> { + validate_index_dtype(indices)?; + let flat = tensor.contiguous().flatten()?; + let indices_flat = indices.contiguous().flatten()?; + let out_flat = self.index_select(&flat, 0, &indices_flat)?; + out_flat.reshape(indices.shape()) + } + + /// Put values into a tensor at flat indices (functional, non-mutating). + /// + /// Returns a new tensor with `values` written at positions specified by `indices`, + /// treating the input tensor as flattened 1D storage. + /// + /// # Arguments + /// + /// * `tensor` - Input tensor to update + /// * `indices` - Index tensor (I32 or I64) containing flat indices + /// * `values` - Values to write. Must have the same number of elements as `indices`. + /// + /// # Returns + /// + /// New tensor with the same shape as `tensor` and updated values + fn put( + &self, + tensor: &Tensor, + indices: &Tensor, + values: &Tensor, + ) -> Result> { + validate_index_dtype(indices)?; + let flat = tensor.contiguous().flatten()?; + let indices_flat = indices.contiguous().flatten()?; + let values_flat = values.contiguous().flatten()?; + + if values_flat.numel() != indices_flat.numel() { + return Err(Error::ShapeMismatch { + expected: vec![indices_flat.numel()], + got: vec![values_flat.numel()], + }); + } + + let out_flat = self.index_put(&flat, 0, &indices_flat, &values_flat)?; + out_flat.reshape(tensor.shape()) + } + /// Select elements where mask is true, returning a flattened 1D tensor. /// /// # Arguments @@ -256,7 +328,7 @@ pub trait IndexingOps { /// /// * `embeddings` - 2D embedding table of shape `` `[vocab_size, embedding_dim]` `` /// * `indices` - Index tensor of any shape containing indices into the embedding table. - /// Must be I64 (or I32 on WebGPU). Values must be in range `` `[0, vocab_size)` ``. + /// Must be I32 or I64. Values must be in range `` `[0, vocab_size)` ``. /// /// # Returns /// @@ -278,7 +350,7 @@ pub trait IndexingOps { /// # Errors /// /// * `ShapeMismatch` - if embeddings is not 2D - /// * `DTypeMismatch` - if indices is not I64 (or I32 on WebGPU) + /// * `DTypeMismatch` - if indices is not I32 or I64 /// * Index out of bounds results in undefined behavior (implementation may return zeros) /// /// # Performance @@ -312,7 +384,7 @@ pub trait IndexingOps { /// /// * `dst` - Destination tensor to scatter into (used as initial values) /// * `dim` - Dimension along which to scatter - /// * `index` - Index tensor (I64) specifying scatter positions + /// * `index` - Index tensor (I32 or I64) specifying scatter positions /// * `src` - Source tensor with values to scatter /// * `op` - Reduction operation to apply (Sum, Mean, Max, Min, Prod) /// * `include_self` - If true, include `dst` values in reduction; if false, initialize @@ -453,8 +525,8 @@ pub trait IndexingOps { /// # Arguments /// /// * `input` - 2D input tensor of shape `` `[nrows, ncols]` `` - /// * `rows` - 1D index tensor (I64) specifying row indices - /// * `cols` - 1D index tensor (I64) specifying column indices + /// * `rows` - 1D index tensor (I32 or I64) specifying row indices + /// * `cols` - 1D index tensor (I32 or I64) specifying column indices /// /// # Returns /// @@ -477,7 +549,7 @@ pub trait IndexingOps { /// # Errors /// /// * `ShapeMismatch` - if input is not 2D or rows/cols have different lengths - /// * `DTypeMismatch` - if rows or cols are not I64 + /// * `DTypeMismatch` - if rows or cols are not I32 or I64 /// * `IndexOutOfBounds` - if any (row, col) pair is out of bounds fn gather_2d( &self, diff --git a/src/ops/wgpu/indexing.rs b/src/ops/wgpu/indexing.rs index c5241fad..372ba88d 100644 --- a/src/ops/wgpu/indexing.rs +++ b/src/ops/wgpu/indexing.rs @@ -2,14 +2,14 @@ use crate::dtype::DType; use crate::error::{Error, Result}; -use crate::ops::{IndexingOps, ReduceOps, ScatterReduceOp, TypeConversionOps}; +use crate::ops::{IndexingOps, ReduceOps, ScalarOps, ScatterReduceOp, TypeConversionOps}; use crate::runtime::RuntimeClient; use crate::runtime::ensure_contiguous; use crate::runtime::wgpu::WgpuClient; use crate::runtime::wgpu::WgpuRuntime; use crate::runtime::wgpu::ops::helpers::{ - BincountParams, Gather2dParams, GatherNdParams, ScatterReduceParams, alloc_output, - create_params_buffer, ensure_i32_indices, get_tensor_buffer, + BincountParams, Gather2dParams, GatherNdParams, MeanDivParams, ScatterReduceParams, + alloc_output, create_params_buffer, ensure_i32_indices, get_tensor_buffer, }; use crate::runtime::wgpu::ops::native::{ native_argreduce_op, native_embedding_lookup, native_gather, native_index_put, @@ -17,6 +17,7 @@ use crate::runtime::wgpu::ops::native::{ }; use crate::runtime::wgpu::shaders::{ launch_bincount, launch_gather_2d, launch_gather_nd, launch_scatter_reduce, + launch_scatter_reduce_count, launch_scatter_reduce_mean_div, launch_scatter_reduce_prod, }; use crate::tensor::Tensor; @@ -122,6 +123,14 @@ impl IndexingOps for WgpuClient { }); } + // Mean only supports F32 on WebGPU + if matches!(op, ScatterReduceOp::Mean) && dtype != DType::F32 { + return Err(Error::UnsupportedDType { + dtype, + op: "scatter_reduce_mean", + }); + } + // Validate index dtype if !matches!(index.dtype(), DType::I32 | DType::I64) { return Err(Error::InvalidArgument { @@ -130,23 +139,6 @@ impl IndexingOps for WgpuClient { }); } - // Map operation - let op_str = match op { - ScatterReduceOp::Sum => "sum", - ScatterReduceOp::Max => "max", - ScatterReduceOp::Min => "min", - ScatterReduceOp::Prod => { - return Err(Error::NotImplemented { - feature: "scatter_reduce with Prod on WebGPU", - }); - } - ScatterReduceOp::Mean => { - return Err(Error::NotImplemented { - feature: "scatter_reduce with Mean on WebGPU (requires count tracking)", - }); - } - }; - // Ensure contiguous let dst = ensure_contiguous(dst); let index_i32 = ensure_i32_indices(self, index)?; @@ -169,26 +161,22 @@ impl IndexingOps for WgpuClient { let src_dim_size = src.shape().get(dim).copied().unwrap_or(1); let total_src = src.numel(); - // Allocate output and initialize + // Initialize output with identity for the operation + let identity = match op { + ScatterReduceOp::Sum | ScatterReduceOp::Mean => 0.0f64, + ScatterReduceOp::Max => f64::NEG_INFINITY, + ScatterReduceOp::Min => f64::INFINITY, + ScatterReduceOp::Prod => 1.0, + }; let output = if include_self { - dst.clone() + // Must deep-copy: clone() shares the GPU buffer, but scatter_reduce + // modifies it in-place via atomics, which would corrupt the original. + self.add_scalar(&dst, 0.0)? } else { - // Initialize to identity for the operation - let identity = match op { - ScatterReduceOp::Sum => 0.0f64, - ScatterReduceOp::Max => f64::NEG_INFINITY, - ScatterReduceOp::Min => f64::INFINITY, - _ => 0.0, - }; Tensor::full_scalar(dst_shape, dtype, identity, self.device()) }; - // Get buffers - let src_buf = get_tensor_buffer(&src)?; - let index_buf = get_tensor_buffer(&index)?; - let output_buf = get_tensor_buffer(&output)?; - - // Create params + // Create shared params let params = ScatterReduceParams { dim: dim as u32, outer_size: outer_size as u32, @@ -201,19 +189,105 @@ impl IndexingOps for WgpuClient { }; let params_buf = create_params_buffer(self, ¶ms); - launch_scatter_reduce( - self.pipeline_cache(), - self.wgpu_queue(), - &src_buf, - &index_buf, - &output_buf, - ¶ms_buf, - total_src, - dtype, - op_str, - )?; + let src_buf = get_tensor_buffer(&src)?; + let index_buf = get_tensor_buffer(&index)?; + let output_buf = get_tensor_buffer(&output)?; - Ok(output) + match op { + ScatterReduceOp::Prod => { + launch_scatter_reduce_prod( + self.pipeline_cache(), + self.wgpu_queue(), + &src_buf, + &index_buf, + &output_buf, + ¶ms_buf, + total_src, + dtype, + )?; + Ok(output) + } + ScatterReduceOp::Mean => { + // Step 1: scatter sum + launch_scatter_reduce( + self.pipeline_cache(), + self.wgpu_queue(), + &src_buf, + &index_buf, + &output_buf, + ¶ms_buf, + total_src, + dtype, + "sum", + )?; + + // Step 2: scatter count (u32 buffer) + let numel = dst.numel(); + let count_init = if include_self { 1u32 } else { 0u32 }; + let count_data = vec![count_init; numel]; + let count_tensor = + Tensor::::from_slice(&count_data, dst_shape, self.device()); + let count_buf = get_tensor_buffer(&count_tensor)?; + + launch_scatter_reduce_count( + self.pipeline_cache(), + self.wgpu_queue(), + &index_buf, + &count_buf, + ¶ms_buf, + total_src, + dtype, + )?; + + // Step 3: divide sum by count + let result = alloc_output(self, dst_shape, dtype); + let result_buf = get_tensor_buffer(&result)?; + + let mean_params = MeanDivParams { + n: numel as u32, + _pad0: 0, + _pad1: 0, + _pad2: 0, + }; + let mean_params_buf = create_params_buffer(self, &mean_params); + + launch_scatter_reduce_mean_div( + self.pipeline_cache(), + self.wgpu_queue(), + &output_buf, + &count_buf, + &result_buf, + &mean_params_buf, + numel, + dtype, + )?; + + Ok(result) + } + _ => { + // Sum, Max, Min - use existing shader + let op_str = match op { + ScatterReduceOp::Sum => "sum", + ScatterReduceOp::Max => "max", + ScatterReduceOp::Min => "min", + _ => unreachable!(), + }; + + launch_scatter_reduce( + self.pipeline_cache(), + self.wgpu_queue(), + &src_buf, + &index_buf, + &output_buf, + ¶ms_buf, + total_src, + dtype, + op_str, + )?; + + Ok(output) + } + } } fn gather_nd( diff --git a/src/runtime/cpu/helpers/indexing.rs b/src/runtime/cpu/helpers/indexing.rs index 4fd35cdd..f0a5355e 100644 --- a/src/runtime/cpu/helpers/indexing.rs +++ b/src/runtime/cpu/helpers/indexing.rs @@ -11,6 +11,32 @@ use crate::ops::ScatterReduceOp; use crate::runtime::ensure_contiguous; use crate::tensor::Tensor; +/// Normalize index tensor to I64 for CPU indexing kernels. +/// +/// CPU indexing kernels consume i64 indices. This helper accepts both I32 and I64 +/// public API inputs and materializes an I64 tensor when needed. +fn normalize_indices_to_i64( + client: &CpuClient, + indices: &Tensor, +) -> Result> { + match indices.dtype() { + DType::I64 => Ok(indices.clone()), + DType::I32 => { + let idx_i32: Vec = indices.to_vec(); + let idx_i64: Vec = idx_i32.into_iter().map(i64::from).collect(); + Ok(Tensor::::from_slice( + &idx_i64, + indices.shape(), + &client.device, + )) + } + other => Err(Error::DTypeMismatch { + lhs: DType::I64, + rhs: other, + }), + } +} + /// Gather elements from a 2D matrix using row and column index vectors. /// /// For each index i, extracts `input[rows[i], cols[i]]`. @@ -34,48 +60,36 @@ pub fn gather_2d_impl( let nrows = shape[0]; let ncols = shape[1]; - // Validate index dtypes - if rows.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: rows.dtype(), - }); - } - - if cols.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: cols.dtype(), - }); - } + let rows_i64 = normalize_indices_to_i64(client, rows)?; + let cols_i64 = normalize_indices_to_i64(client, cols)?; // Validate rows and cols are 1D and have same length - if rows.ndim() != 1 { + if rows_i64.ndim() != 1 { return Err(Error::ShapeMismatch { - expected: vec![rows.numel()], - got: rows.shape().to_vec(), + expected: vec![rows_i64.numel()], + got: rows_i64.shape().to_vec(), }); } - if cols.ndim() != 1 { + if cols_i64.ndim() != 1 { return Err(Error::ShapeMismatch { - expected: vec![cols.numel()], - got: cols.shape().to_vec(), + expected: vec![cols_i64.numel()], + got: cols_i64.shape().to_vec(), }); } - let num_indices = rows.numel(); - if cols.numel() != num_indices { + let num_indices = rows_i64.numel(); + if cols_i64.numel() != num_indices { return Err(Error::ShapeMismatch { expected: vec![num_indices], - got: cols.shape().to_vec(), + got: cols_i64.shape().to_vec(), }); } // Make all inputs contiguous let input_contig = ensure_contiguous(input); - let rows_contig = ensure_contiguous(rows); - let cols_contig = ensure_contiguous(cols); + let rows_contig = ensure_contiguous(&rows_i64); + let cols_contig = ensure_contiguous(&cols_i64); // Allocate output let out = Tensor::::empty(&[num_indices], dtype, &client.device); @@ -127,27 +141,21 @@ pub fn gather_impl( }); } - // Validate index dtype - if index.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: index.dtype(), - }); - } + let index_i64 = normalize_indices_to_i64(client, index)?; // Validate index dimensions - if index.ndim() != ndim { + if index_i64.ndim() != ndim { return Err(Error::ShapeMismatch { expected: shape.to_vec(), - got: index.shape().to_vec(), + got: index_i64.shape().to_vec(), }); } // Output shape is same as index shape - let out_shape = index.shape().to_vec(); + let out_shape = index_i64.shape().to_vec(); let a_contig = ensure_contiguous(a); - let index_contig = ensure_contiguous(index); + let index_contig = ensure_contiguous(&index_i64); let out = Tensor::::empty(&out_shape, dtype, &client.device); let a_ptr = a_contig.storage().ptr(); @@ -190,13 +198,7 @@ pub fn scatter_impl( }); } - // Validate dtypes - if index.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: index.dtype(), - }); - } + let index_i64 = normalize_indices_to_i64(client, index)?; if src.dtype() != dtype { return Err(Error::DTypeMismatch { @@ -206,15 +208,15 @@ pub fn scatter_impl( } // Validate shapes - if index.shape() != src.shape() { + if index_i64.shape() != src.shape() { return Err(Error::ShapeMismatch { expected: src.shape().to_vec(), - got: index.shape().to_vec(), + got: index_i64.shape().to_vec(), }); } let a_contig = ensure_contiguous(a); - let index_contig = ensure_contiguous(index); + let index_contig = ensure_contiguous(&index_i64); let src_contig = ensure_contiguous(src); let out = Tensor::::empty(shape, dtype, &client.device); @@ -231,7 +233,7 @@ pub fn scatter_impl( src_ptr as *const T, out_ptr as *mut T, shape, - index.shape(), + index_i64.shape(), dim, ); } @@ -259,30 +261,24 @@ pub fn index_select_impl( }); } - // Validate index dtype - if index.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: index.dtype(), - }); - } + let index_i64 = normalize_indices_to_i64(client, index)?; // Index must be 1D - if index.ndim() != 1 { + if index_i64.ndim() != 1 { return Err(Error::ShapeMismatch { - expected: vec![index.numel()], - got: index.shape().to_vec(), + expected: vec![index_i64.numel()], + got: index_i64.shape().to_vec(), }); } - let index_len = index.shape()[0]; + let index_len = index_i64.shape()[0]; // Output shape: replace dimension `dim` with index length let mut out_shape = shape.to_vec(); out_shape[dim] = index_len; let a_contig = ensure_contiguous(a); - let index_contig = ensure_contiguous(index); + let index_contig = ensure_contiguous(&index_i64); // Validate all indices are within bounds (before calling unsafe kernel) let dim_size = shape[dim]; @@ -341,19 +337,13 @@ pub fn index_put_impl( }); } - // Validate index dtype - if index.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: index.dtype(), - }); - } + let index_i64 = normalize_indices_to_i64(client, index)?; // Index must be 1D - if index.ndim() != 1 { + if index_i64.ndim() != 1 { return Err(Error::ShapeMismatch { - expected: vec![index.numel()], - got: index.shape().to_vec(), + expected: vec![index_i64.numel()], + got: index_i64.shape().to_vec(), }); } @@ -365,7 +355,7 @@ pub fn index_put_impl( }); } - let index_len = index.shape()[0]; + let index_len = index_i64.shape()[0]; // Validate src shape: must match a's shape except at dim where it equals index_len let mut expected_src_shape = shape.to_vec(); @@ -378,7 +368,7 @@ pub fn index_put_impl( } let a_contig = ensure_contiguous(a); - let index_contig = ensure_contiguous(index); + let index_contig = ensure_contiguous(&index_i64); let src_contig = ensure_contiguous(src); // Validate all indices are within bounds (before calling unsafe kernel) @@ -622,24 +612,18 @@ pub fn embedding_lookup_impl( }); } - // Validate indices dtype - if indices.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: indices.dtype(), - }); - } + let indices_i64 = normalize_indices_to_i64(client, indices)?; let vocab_size = emb_shape[0]; let embedding_dim = emb_shape[1]; - let num_indices = indices.numel(); + let num_indices = indices_i64.numel(); // Output shape: indices.shape() + [embedding_dim] - let mut out_shape = indices.shape().to_vec(); + let mut out_shape = indices_i64.shape().to_vec(); out_shape.push(embedding_dim); let emb_contig = ensure_contiguous(embeddings); - let idx_contig = ensure_contiguous(indices); + let idx_contig = ensure_contiguous(&indices_i64); let out = Tensor::::empty(&out_shape, dtype, &client.device); let emb_ptr = emb_contig.storage().ptr(); @@ -685,13 +669,7 @@ pub fn scatter_reduce_impl( }); } - // Validate dtypes - if index.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: index.dtype(), - }); - } + let index_i64 = normalize_indices_to_i64(client, index)?; if src.dtype() != dtype { return Err(Error::DTypeMismatch { @@ -701,23 +679,23 @@ pub fn scatter_reduce_impl( } // Validate that index and src have same shape - if index.shape() != src.shape() { + if index_i64.shape() != src.shape() { return Err(Error::ShapeMismatch { expected: src.shape().to_vec(), - got: index.shape().to_vec(), + got: index_i64.shape().to_vec(), }); } // Validate that index has same number of dimensions as dst - if index.ndim() != ndim { + if index_i64.ndim() != ndim { return Err(Error::ShapeMismatch { expected: shape.to_vec(), - got: index.shape().to_vec(), + got: index_i64.shape().to_vec(), }); } let dst_contig = ensure_contiguous(dst); - let index_contig = ensure_contiguous(index); + let index_contig = ensure_contiguous(&index_i64); let src_contig = ensure_contiguous(src); let out = Tensor::::empty(shape, dtype, &client.device); @@ -744,7 +722,7 @@ pub fn scatter_reduce_impl( out_ptr as *mut T, counts_ptr, shape, - index.shape(), + index_i64.shape(), dim, op, include_self, @@ -763,15 +741,8 @@ pub fn gather_nd_impl( ) -> Result> { let dtype = input.dtype(); let input_shape = input.shape(); - let indices_shape = indices.shape(); - - // Validate indices dtype - if indices.dtype() != DType::I64 { - return Err(Error::DTypeMismatch { - lhs: DType::I64, - rhs: indices.dtype(), - }); - } + let indices_i64 = normalize_indices_to_i64(client, indices)?; + let indices_shape = indices_i64.shape(); // Indices must have at least 1 dimension if indices_shape.is_empty() { @@ -804,7 +775,7 @@ pub fn gather_nd_impl( } let input_contig = ensure_contiguous(input); - let indices_contig = ensure_contiguous(indices); + let indices_contig = ensure_contiguous(&indices_i64); let out = Tensor::::empty(&out_shape, dtype, &client.device); let input_ptr = input_contig.storage().ptr(); diff --git a/src/runtime/cpu/kernels/complex.rs b/src/runtime/cpu/kernels/complex.rs index a1eee01a..2bc8525d 100644 --- a/src/runtime/cpu/kernels/complex.rs +++ b/src/runtime/cpu/kernels/complex.rs @@ -17,23 +17,29 @@ use crate::dtype::{Complex64, Complex128}; use rayon::prelude::*; /// Parallelization threshold: skip Rayon for small tensors (overhead > benefit) +#[cfg(feature = "rayon")] const PARALLEL_THRESHOLD: usize = 4096; /// Complex conjugate kernel for Complex64 with Rayon parallelization /// /// Performance: ~95% of memory bandwidth (memory-bound, simple negation) #[inline] -pub unsafe fn conj_complex64(input: *const Complex64, output: *mut Complex64, numel: usize) { +pub unsafe fn conj_complex64( + input: *const Complex64, + output: *mut Complex64, + numel: usize, + _chunk_size: usize, +) { #[cfg(feature = "rayon")] if numel >= PARALLEL_THRESHOLD { use std::slice; let input_slice = slice::from_raw_parts(input, numel); let output_slice = slice::from_raw_parts_mut(output, numel); + let chunk_size = _chunk_size.max(1); - const CHUNK_SIZE: usize = 4096; output_slice - .par_chunks_mut(CHUNK_SIZE) - .zip(input_slice.par_chunks(CHUNK_SIZE)) + .par_chunks_mut(chunk_size) + .zip(input_slice.par_chunks(chunk_size)) .for_each(|(out_chunk, in_chunk)| { for (o, &i) in out_chunk.iter_mut().zip(in_chunk) { *o = i.conj(); @@ -50,17 +56,22 @@ pub unsafe fn conj_complex64(input: *const Complex64, output: *mut Complex64, nu /// Complex conjugate kernel for Complex128 with Rayon parallelization #[inline] -pub unsafe fn conj_complex128(input: *const Complex128, output: *mut Complex128, numel: usize) { +pub unsafe fn conj_complex128( + input: *const Complex128, + output: *mut Complex128, + numel: usize, + _chunk_size: usize, +) { #[cfg(feature = "rayon")] if numel >= PARALLEL_THRESHOLD { use std::slice; let input_slice = slice::from_raw_parts(input, numel); let output_slice = slice::from_raw_parts_mut(output, numel); + let chunk_size = _chunk_size.max(1); - const CHUNK_SIZE: usize = 4096; output_slice - .par_chunks_mut(CHUNK_SIZE) - .zip(input_slice.par_chunks(CHUNK_SIZE)) + .par_chunks_mut(chunk_size) + .zip(input_slice.par_chunks(chunk_size)) .for_each(|(out_chunk, in_chunk)| { for (o, &i) in out_chunk.iter_mut().zip(in_chunk) { *o = i.conj(); @@ -78,17 +89,22 @@ pub unsafe fn conj_complex128(input: *const Complex128, output: *mut Complex128, /// /// Performance: ~98% of memory bandwidth (pure memory copy) #[inline] -pub unsafe fn real_complex64(input: *const Complex64, output: *mut f32, numel: usize) { +pub unsafe fn real_complex64( + input: *const Complex64, + output: *mut f32, + numel: usize, + _chunk_size: usize, +) { #[cfg(feature = "rayon")] if numel >= PARALLEL_THRESHOLD { use std::slice; let input_slice = slice::from_raw_parts(input, numel); let output_slice = slice::from_raw_parts_mut(output, numel); + let chunk_size = _chunk_size.max(1); - const CHUNK_SIZE: usize = 4096; output_slice - .par_chunks_mut(CHUNK_SIZE) - .zip(input_slice.par_chunks(CHUNK_SIZE)) + .par_chunks_mut(chunk_size) + .zip(input_slice.par_chunks(chunk_size)) .for_each(|(out_chunk, in_chunk)| { for (o, &i) in out_chunk.iter_mut().zip(in_chunk) { *o = i.re; @@ -104,17 +120,22 @@ pub unsafe fn real_complex64(input: *const Complex64, output: *mut f32, numel: u /// Extract real component from Complex128 with Rayon parallelization #[inline] -pub unsafe fn real_complex128(input: *const Complex128, output: *mut f64, numel: usize) { +pub unsafe fn real_complex128( + input: *const Complex128, + output: *mut f64, + numel: usize, + _chunk_size: usize, +) { #[cfg(feature = "rayon")] if numel >= PARALLEL_THRESHOLD { use std::slice; let input_slice = slice::from_raw_parts(input, numel); let output_slice = slice::from_raw_parts_mut(output, numel); + let chunk_size = _chunk_size.max(1); - const CHUNK_SIZE: usize = 4096; output_slice - .par_chunks_mut(CHUNK_SIZE) - .zip(input_slice.par_chunks(CHUNK_SIZE)) + .par_chunks_mut(chunk_size) + .zip(input_slice.par_chunks(chunk_size)) .for_each(|(out_chunk, in_chunk)| { for (o, &i) in out_chunk.iter_mut().zip(in_chunk) { *o = i.re; @@ -130,17 +151,22 @@ pub unsafe fn real_complex128(input: *const Complex128, output: *mut f64, numel: /// Extract imaginary component from Complex64 with Rayon parallelization #[inline] -pub unsafe fn imag_complex64(input: *const Complex64, output: *mut f32, numel: usize) { +pub unsafe fn imag_complex64( + input: *const Complex64, + output: *mut f32, + numel: usize, + _chunk_size: usize, +) { #[cfg(feature = "rayon")] if numel >= PARALLEL_THRESHOLD { use std::slice; let input_slice = slice::from_raw_parts(input, numel); let output_slice = slice::from_raw_parts_mut(output, numel); + let chunk_size = _chunk_size.max(1); - const CHUNK_SIZE: usize = 4096; output_slice - .par_chunks_mut(CHUNK_SIZE) - .zip(input_slice.par_chunks(CHUNK_SIZE)) + .par_chunks_mut(chunk_size) + .zip(input_slice.par_chunks(chunk_size)) .for_each(|(out_chunk, in_chunk)| { for (o, &i) in out_chunk.iter_mut().zip(in_chunk) { *o = i.im; @@ -156,17 +182,22 @@ pub unsafe fn imag_complex64(input: *const Complex64, output: *mut f32, numel: u /// Extract imaginary component from Complex128 with Rayon parallelization #[inline] -pub unsafe fn imag_complex128(input: *const Complex128, output: *mut f64, numel: usize) { +pub unsafe fn imag_complex128( + input: *const Complex128, + output: *mut f64, + numel: usize, + _chunk_size: usize, +) { #[cfg(feature = "rayon")] if numel >= PARALLEL_THRESHOLD { use std::slice; let input_slice = slice::from_raw_parts(input, numel); let output_slice = slice::from_raw_parts_mut(output, numel); + let chunk_size = _chunk_size.max(1); - const CHUNK_SIZE: usize = 4096; output_slice - .par_chunks_mut(CHUNK_SIZE) - .zip(input_slice.par_chunks(CHUNK_SIZE)) + .par_chunks_mut(chunk_size) + .zip(input_slice.par_chunks(chunk_size)) .for_each(|(out_chunk, in_chunk)| { for (o, &i) in out_chunk.iter_mut().zip(in_chunk) { *o = i.im; @@ -184,17 +215,22 @@ pub unsafe fn imag_complex128(input: *const Complex128, output: *mut f64, numel: /// /// Performance: ~80% of compute bound (atan2 ~20 cycles) #[inline] -pub unsafe fn angle_complex64(input: *const Complex64, output: *mut f32, numel: usize) { +pub unsafe fn angle_complex64( + input: *const Complex64, + output: *mut f32, + numel: usize, + _chunk_size: usize, +) { #[cfg(feature = "rayon")] if numel >= PARALLEL_THRESHOLD { use std::slice; let input_slice = slice::from_raw_parts(input, numel); let output_slice = slice::from_raw_parts_mut(output, numel); + let chunk_size = _chunk_size.max(1); - const CHUNK_SIZE: usize = 4096; output_slice - .par_chunks_mut(CHUNK_SIZE) - .zip(input_slice.par_chunks(CHUNK_SIZE)) + .par_chunks_mut(chunk_size) + .zip(input_slice.par_chunks(chunk_size)) .for_each(|(out_chunk, in_chunk)| { for (o, &i) in out_chunk.iter_mut().zip(in_chunk) { *o = i.phase(); @@ -210,17 +246,22 @@ pub unsafe fn angle_complex64(input: *const Complex64, output: *mut f32, numel: /// Compute phase angle for Complex128 with Rayon parallelization #[inline] -pub unsafe fn angle_complex128(input: *const Complex128, output: *mut f64, numel: usize) { +pub unsafe fn angle_complex128( + input: *const Complex128, + output: *mut f64, + numel: usize, + _chunk_size: usize, +) { #[cfg(feature = "rayon")] if numel >= PARALLEL_THRESHOLD { use std::slice; let input_slice = slice::from_raw_parts(input, numel); let output_slice = slice::from_raw_parts_mut(output, numel); + let chunk_size = _chunk_size.max(1); - const CHUNK_SIZE: usize = 4096; output_slice - .par_chunks_mut(CHUNK_SIZE) - .zip(input_slice.par_chunks(CHUNK_SIZE)) + .par_chunks_mut(chunk_size) + .zip(input_slice.par_chunks(chunk_size)) .for_each(|(out_chunk, in_chunk)| { for (o, &i) in out_chunk.iter_mut().zip(in_chunk) { *o = i.phase(); @@ -238,17 +279,22 @@ pub unsafe fn angle_complex128(input: *const Complex128, output: *mut f64, numel /// /// angle(x) = 0 if x >= 0, π if x < 0 #[inline] -pub unsafe fn angle_real_f32(input: *const f32, output: *mut f32, numel: usize) { +pub unsafe fn angle_real_f32( + input: *const f32, + output: *mut f32, + numel: usize, + _chunk_size: usize, +) { #[cfg(feature = "rayon")] if numel >= PARALLEL_THRESHOLD { use std::slice; let input_slice = slice::from_raw_parts(input, numel); let output_slice = slice::from_raw_parts_mut(output, numel); + let chunk_size = _chunk_size.max(1); - const CHUNK_SIZE: usize = 4096; output_slice - .par_chunks_mut(CHUNK_SIZE) - .zip(input_slice.par_chunks(CHUNK_SIZE)) + .par_chunks_mut(chunk_size) + .zip(input_slice.par_chunks(chunk_size)) .for_each(|(out_chunk, in_chunk)| { for (o, &val) in out_chunk.iter_mut().zip(in_chunk) { *o = if val < 0.0 { std::f32::consts::PI } else { 0.0 }; @@ -267,17 +313,22 @@ pub unsafe fn angle_real_f32(input: *const f32, output: *mut f32, numel: usize) /// /// angle(x) = 0 if x >= 0, π if x < 0 #[inline] -pub unsafe fn angle_real_f64(input: *const f64, output: *mut f64, numel: usize) { +pub unsafe fn angle_real_f64( + input: *const f64, + output: *mut f64, + numel: usize, + _chunk_size: usize, +) { #[cfg(feature = "rayon")] if numel >= PARALLEL_THRESHOLD { use std::slice; let input_slice = slice::from_raw_parts(input, numel); let output_slice = slice::from_raw_parts_mut(output, numel); + let chunk_size = _chunk_size.max(1); - const CHUNK_SIZE: usize = 4096; output_slice - .par_chunks_mut(CHUNK_SIZE) - .zip(input_slice.par_chunks(CHUNK_SIZE)) + .par_chunks_mut(chunk_size) + .zip(input_slice.par_chunks(chunk_size)) .for_each(|(out_chunk, in_chunk)| { for (o, &val) in out_chunk.iter_mut().zip(in_chunk) { *o = if val < 0.0 { std::f64::consts::PI } else { 0.0 }; @@ -305,6 +356,7 @@ pub unsafe fn from_real_imag_f32( imag: *const f32, output: *mut Complex64, numel: usize, + _chunk_size: usize, ) { #[cfg(feature = "rayon")] if numel >= PARALLEL_THRESHOLD { @@ -312,12 +364,12 @@ pub unsafe fn from_real_imag_f32( let real_slice = slice::from_raw_parts(real, numel); let imag_slice = slice::from_raw_parts(imag, numel); let output_slice = slice::from_raw_parts_mut(output, numel); + let chunk_size = _chunk_size.max(1); - const CHUNK_SIZE: usize = 4096; output_slice - .par_chunks_mut(CHUNK_SIZE) - .zip(real_slice.par_chunks(CHUNK_SIZE)) - .zip(imag_slice.par_chunks(CHUNK_SIZE)) + .par_chunks_mut(chunk_size) + .zip(real_slice.par_chunks(chunk_size)) + .zip(imag_slice.par_chunks(chunk_size)) .for_each(|((out_chunk, re_chunk), im_chunk)| { for ((o, &re), &im) in out_chunk.iter_mut().zip(re_chunk).zip(im_chunk) { *o = Complex64::new(re, im); @@ -339,6 +391,7 @@ pub unsafe fn from_real_imag_f64( imag: *const f64, output: *mut Complex128, numel: usize, + _chunk_size: usize, ) { #[cfg(feature = "rayon")] if numel >= PARALLEL_THRESHOLD { @@ -346,12 +399,12 @@ pub unsafe fn from_real_imag_f64( let real_slice = slice::from_raw_parts(real, numel); let imag_slice = slice::from_raw_parts(imag, numel); let output_slice = slice::from_raw_parts_mut(output, numel); + let chunk_size = _chunk_size.max(1); - const CHUNK_SIZE: usize = 4096; output_slice - .par_chunks_mut(CHUNK_SIZE) - .zip(real_slice.par_chunks(CHUNK_SIZE)) - .zip(imag_slice.par_chunks(CHUNK_SIZE)) + .par_chunks_mut(chunk_size) + .zip(real_slice.par_chunks(chunk_size)) + .zip(imag_slice.par_chunks(chunk_size)) .for_each(|((out_chunk, re_chunk), im_chunk)| { for ((o, &re), &im) in out_chunk.iter_mut().zip(re_chunk).zip(im_chunk) { *o = Complex128::new(re, im); @@ -378,6 +431,7 @@ pub unsafe fn complex64_mul_real( real: *const f32, output: *mut Complex64, numel: usize, + _chunk_size: usize, ) { #[cfg(feature = "rayon")] if numel >= PARALLEL_THRESHOLD { @@ -385,12 +439,12 @@ pub unsafe fn complex64_mul_real( let complex_slice = slice::from_raw_parts(complex, numel); let real_slice = slice::from_raw_parts(real, numel); let output_slice = slice::from_raw_parts_mut(output, numel); + let chunk_size = _chunk_size.max(1); - const CHUNK_SIZE: usize = 4096; output_slice - .par_chunks_mut(CHUNK_SIZE) - .zip(complex_slice.par_chunks(CHUNK_SIZE)) - .zip(real_slice.par_chunks(CHUNK_SIZE)) + .par_chunks_mut(chunk_size) + .zip(complex_slice.par_chunks(chunk_size)) + .zip(real_slice.par_chunks(chunk_size)) .for_each(|((out_chunk, c_chunk), r_chunk)| { for ((o, &c), &r) in out_chunk.iter_mut().zip(c_chunk).zip(r_chunk) { *o = Complex64::new(c.re * r, c.im * r); @@ -413,6 +467,7 @@ pub unsafe fn complex128_mul_real( real: *const f64, output: *mut Complex128, numel: usize, + _chunk_size: usize, ) { #[cfg(feature = "rayon")] if numel >= PARALLEL_THRESHOLD { @@ -420,12 +475,12 @@ pub unsafe fn complex128_mul_real( let complex_slice = slice::from_raw_parts(complex, numel); let real_slice = slice::from_raw_parts(real, numel); let output_slice = slice::from_raw_parts_mut(output, numel); + let chunk_size = _chunk_size.max(1); - const CHUNK_SIZE: usize = 4096; output_slice - .par_chunks_mut(CHUNK_SIZE) - .zip(complex_slice.par_chunks(CHUNK_SIZE)) - .zip(real_slice.par_chunks(CHUNK_SIZE)) + .par_chunks_mut(chunk_size) + .zip(complex_slice.par_chunks(chunk_size)) + .zip(real_slice.par_chunks(chunk_size)) .for_each(|((out_chunk, c_chunk), r_chunk)| { for ((o, &c), &r) in out_chunk.iter_mut().zip(c_chunk).zip(r_chunk) { *o = Complex128::new(c.re * r, c.im * r); @@ -450,6 +505,7 @@ pub unsafe fn complex64_div_real( real: *const f32, output: *mut Complex64, numel: usize, + _chunk_size: usize, ) { #[cfg(feature = "rayon")] if numel >= PARALLEL_THRESHOLD { @@ -457,12 +513,12 @@ pub unsafe fn complex64_div_real( let complex_slice = slice::from_raw_parts(complex, numel); let real_slice = slice::from_raw_parts(real, numel); let output_slice = slice::from_raw_parts_mut(output, numel); + let chunk_size = _chunk_size.max(1); - const CHUNK_SIZE: usize = 4096; output_slice - .par_chunks_mut(CHUNK_SIZE) - .zip(complex_slice.par_chunks(CHUNK_SIZE)) - .zip(real_slice.par_chunks(CHUNK_SIZE)) + .par_chunks_mut(chunk_size) + .zip(complex_slice.par_chunks(chunk_size)) + .zip(real_slice.par_chunks(chunk_size)) .for_each(|((out_chunk, c_chunk), r_chunk)| { for ((o, &c), &r) in out_chunk.iter_mut().zip(c_chunk).zip(r_chunk) { *o = Complex64::new(c.re / r, c.im / r); @@ -485,6 +541,7 @@ pub unsafe fn complex128_div_real( real: *const f64, output: *mut Complex128, numel: usize, + _chunk_size: usize, ) { #[cfg(feature = "rayon")] if numel >= PARALLEL_THRESHOLD { @@ -492,12 +549,12 @@ pub unsafe fn complex128_div_real( let complex_slice = slice::from_raw_parts(complex, numel); let real_slice = slice::from_raw_parts(real, numel); let output_slice = slice::from_raw_parts_mut(output, numel); + let chunk_size = _chunk_size.max(1); - const CHUNK_SIZE: usize = 4096; output_slice - .par_chunks_mut(CHUNK_SIZE) - .zip(complex_slice.par_chunks(CHUNK_SIZE)) - .zip(real_slice.par_chunks(CHUNK_SIZE)) + .par_chunks_mut(chunk_size) + .zip(complex_slice.par_chunks(chunk_size)) + .zip(real_slice.par_chunks(chunk_size)) .for_each(|((out_chunk, c_chunk), r_chunk)| { for ((o, &c), &r) in out_chunk.iter_mut().zip(c_chunk).zip(r_chunk) { *o = Complex128::new(c.re / r, c.im / r); diff --git a/src/runtime/cpu/kernels/simd/matmul/mod.rs b/src/runtime/cpu/kernels/simd/matmul/mod.rs index ac3b223f..7b15bd97 100644 --- a/src/runtime/cpu/kernels/simd/matmul/mod.rs +++ b/src/runtime/cpu/kernels/simd/matmul/mod.rs @@ -44,6 +44,7 @@ mod tiling; #[cfg(target_arch = "aarch64")] mod aarch64; +#[cfg(all(feature = "f16", target_arch = "x86_64"))] pub(crate) mod half_convert; use super::{SimdLevel, detect_simd}; diff --git a/src/runtime/cuda/kernels/index.cu b/src/runtime/cuda/kernels/index.cu index 553818e3..7e7d4933 100644 --- a/src/runtime/cuda/kernels/index.cu +++ b/src/runtime/cuda/kernels/index.cu @@ -1058,4 +1058,139 @@ DEFINE_GATHER_2D_KERNEL(bf16, __nv_bfloat16) DEFINE_GATHER_2D_KERNEL(i32, int) DEFINE_GATHER_2D_KERNEL(i64, long long) +// ============================================================================ +// Scatter Reduce - Prod (atomic multiply via CAS) +// ============================================================================ + +__device__ __forceinline__ float atomicMulFloat(float* address, float val) { + int* address_as_int = (int*)address; + int old = *address_as_int; + int assumed; + do { + assumed = old; + float old_val = __int_as_float(assumed); + float new_val = old_val * val; + old = atomicCAS(address_as_int, assumed, __float_as_int(new_val)); + } while (assumed != old); + return __int_as_float(old); +} + +__device__ __forceinline__ double atomicMulDouble(double* address, double val) { + unsigned long long* address_as_ull = (unsigned long long*)address; + unsigned long long old = *address_as_ull; + unsigned long long assumed; + do { + assumed = old; + double old_val = __longlong_as_double(assumed); + double new_val = old_val * val; + old = atomicCAS(address_as_ull, assumed, __double_as_longlong(new_val)); + } while (assumed != old); + return __longlong_as_double(old); +} + +__device__ __forceinline__ int atomicMulInt(int* address, int val) { + int old = *address; + int assumed; + do { + assumed = old; + int new_val = assumed * val; + old = atomicCAS(address, assumed, new_val); + } while (assumed != old); + return old; +} + +#define DEFINE_SCATTER_REDUCE_PROD_KERNEL(suffix, dtype, atomic_mul_fn) \ +__global__ void scatter_reduce_prod_##suffix( \ + const dtype* __restrict__ src, \ + const long long* __restrict__ indices, \ + dtype* __restrict__ dst, \ + unsigned int dim, \ + unsigned int outer_size, \ + unsigned int dim_size, \ + unsigned int inner_size, \ + unsigned int src_dim_size \ +) { \ + unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x; \ + unsigned int total = outer_size * src_dim_size * inner_size; \ + if (idx >= total) return; \ + \ + unsigned int inner = idx % inner_size; \ + unsigned int src_d = (idx / inner_size) % src_dim_size; \ + unsigned int outer = idx / (src_dim_size * inner_size); \ + \ + long long index_val = indices[src_d]; \ + if (index_val < 0 || (unsigned int)index_val >= dim_size) { \ + return; \ + } \ + \ + unsigned int dst_idx = outer * dim_size * inner_size + (unsigned int)index_val * inner_size + inner; \ + unsigned int src_idx = outer * src_dim_size * inner_size + src_d * inner_size + inner; \ + \ + atomic_mul_fn(&dst[dst_idx], src[src_idx]); \ +} + +DEFINE_SCATTER_REDUCE_PROD_KERNEL(f32, float, atomicMulFloat) +DEFINE_SCATTER_REDUCE_PROD_KERNEL(f64, double, atomicMulDouble) +DEFINE_SCATTER_REDUCE_PROD_KERNEL(i32, int, atomicMulInt) + +// ============================================================================ +// Scatter Reduce - Count (for mean: atomicAdd 1 to count buffer per scatter) +// ============================================================================ + +#define DEFINE_SCATTER_REDUCE_COUNT_KERNEL(suffix, dtype) \ +__global__ void scatter_reduce_count_##suffix( \ + const long long* __restrict__ indices, \ + dtype* __restrict__ count, \ + unsigned int dim, \ + unsigned int outer_size, \ + unsigned int dim_size, \ + unsigned int inner_size, \ + unsigned int src_dim_size \ +) { \ + unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x; \ + unsigned int total = outer_size * src_dim_size * inner_size; \ + if (idx >= total) return; \ + \ + unsigned int inner = idx % inner_size; \ + unsigned int src_d = (idx / inner_size) % src_dim_size; \ + unsigned int outer = idx / (src_dim_size * inner_size); \ + \ + long long index_val = indices[src_d]; \ + if (index_val < 0 || (unsigned int)index_val >= dim_size) { \ + return; \ + } \ + \ + unsigned int dst_idx = outer * dim_size * inner_size + (unsigned int)index_val * inner_size + inner; \ + \ + atomicAdd(&count[dst_idx], (dtype)1); \ +} + +DEFINE_SCATTER_REDUCE_COUNT_KERNEL(f32, float) +DEFINE_SCATTER_REDUCE_COUNT_KERNEL(f64, double) + +// ============================================================================ +// Scatter Reduce - Mean Divide (element-wise: out[i] = sum[i] / count[i]) +// ============================================================================ + +#define DEFINE_SCATTER_REDUCE_MEAN_DIV_KERNEL(suffix, dtype) \ +__global__ void scatter_reduce_mean_div_##suffix( \ + const dtype* __restrict__ sum_buf, \ + const dtype* __restrict__ count_buf, \ + dtype* __restrict__ output, \ + unsigned int n \ +) { \ + unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x; \ + if (idx >= n) return; \ + \ + dtype c = count_buf[idx]; \ + if (c > (dtype)0) { \ + output[idx] = sum_buf[idx] / c; \ + } else { \ + output[idx] = (dtype)0; \ + } \ +} + +DEFINE_SCATTER_REDUCE_MEAN_DIV_KERNEL(f32, float) +DEFINE_SCATTER_REDUCE_MEAN_DIV_KERNEL(f64, double) + } // extern "C" diff --git a/src/runtime/cuda/kernels/index.rs b/src/runtime/cuda/kernels/index.rs index a62d3271..ecd06924 100644 --- a/src/runtime/cuda/kernels/index.rs +++ b/src/runtime/cuda/kernels/index.rs @@ -1134,6 +1134,7 @@ pub enum ScatterReduceOpCuda { Sum, Max, Min, + Prod, } /// Launch scatter_reduce kernel. @@ -1179,12 +1180,15 @@ pub unsafe fn launch_scatter_reduce( (DType::F32, ScatterReduceOpCuda::Sum) => "scatter_reduce_sum_f32", (DType::F32, ScatterReduceOpCuda::Max) => "scatter_reduce_max_f32", (DType::F32, ScatterReduceOpCuda::Min) => "scatter_reduce_min_f32", + (DType::F32, ScatterReduceOpCuda::Prod) => "scatter_reduce_prod_f32", (DType::F64, ScatterReduceOpCuda::Sum) => "scatter_reduce_sum_f64", (DType::F64, ScatterReduceOpCuda::Max) => "scatter_reduce_max_f64", (DType::F64, ScatterReduceOpCuda::Min) => "scatter_reduce_min_f64", + (DType::F64, ScatterReduceOpCuda::Prod) => "scatter_reduce_prod_f64", (DType::I32, ScatterReduceOpCuda::Sum) => "scatter_reduce_sum_i32", (DType::I32, ScatterReduceOpCuda::Max) => "scatter_reduce_max_i32", (DType::I32, ScatterReduceOpCuda::Min) => "scatter_reduce_min_i32", + (DType::I32, ScatterReduceOpCuda::Prod) => "scatter_reduce_prod_i32", _ => { return Err(Error::UnsupportedDType { dtype, @@ -1223,6 +1227,149 @@ pub unsafe fn launch_scatter_reduce( } } +// ============================================================================ +// Scatter Reduce Count (for mean) +// ============================================================================ + +/// Launch scatter_reduce_count kernel. +/// +/// Atomically increments count buffer at scattered positions. +/// Used as part of scatter_reduce mean: sum / count. +/// +/// # Safety +/// +/// All pointers must be valid device memory. +#[allow(clippy::too_many_arguments)] +pub unsafe fn launch_scatter_reduce_count( + context: &Arc, + stream: &CudaStream, + device_index: usize, + dtype: DType, + indices_ptr: u64, + count_ptr: u64, + dim: usize, + outer_size: usize, + dim_size: usize, + inner_size: usize, + src_dim_size: usize, +) -> Result<()> { + let total = outer_size * src_dim_size * inner_size; + if total == 0 { + return Ok(()); + } + + unsafe { + let module = get_or_load_module(context, device_index, INDEX_MODULE)?; + + let func_name = match dtype { + DType::F32 => "scatter_reduce_count_f32", + DType::F64 => "scatter_reduce_count_f64", + _ => { + return Err(Error::UnsupportedDType { + dtype, + op: "scatter_reduce_count", + }); + } + }; + + let func = get_kernel_function(&module, func_name)?; + + let grid = elementwise_launch_config(total); + let block = (BLOCK_SIZE, 1, 1); + let cfg = launch_config(grid, block, 0); + + let dim_u32 = dim as u32; + let outer_size_u32 = outer_size as u32; + let dim_size_u32 = dim_size as u32; + let inner_size_u32 = inner_size as u32; + let src_dim_size_u32 = src_dim_size as u32; + + let mut builder = stream.launch_builder(&func); + builder.arg(&indices_ptr); + builder.arg(&count_ptr); + builder.arg(&dim_u32); + builder.arg(&outer_size_u32); + builder.arg(&dim_size_u32); + builder.arg(&inner_size_u32); + builder.arg(&src_dim_size_u32); + + builder.launch(cfg).map_err(|e| { + Error::Internal(format!( + "CUDA scatter_reduce_count kernel launch failed: {:?}", + e + )) + })?; + + Ok(()) + } +} + +// ============================================================================ +// Scatter Reduce Mean Divide +// ============================================================================ + +/// Launch scatter_reduce_mean_div kernel. +/// +/// Element-wise: output[i] = sum[i] / count[i]. +/// If count[i] == 0, output[i] = 0. +/// +/// # Safety +/// +/// All pointers must be valid device memory. +#[allow(clippy::too_many_arguments)] +pub unsafe fn launch_scatter_reduce_mean_div( + context: &Arc, + stream: &CudaStream, + device_index: usize, + dtype: DType, + sum_ptr: u64, + count_ptr: u64, + output_ptr: u64, + n: usize, +) -> Result<()> { + if n == 0 { + return Ok(()); + } + + unsafe { + let module = get_or_load_module(context, device_index, INDEX_MODULE)?; + + let func_name = match dtype { + DType::F32 => "scatter_reduce_mean_div_f32", + DType::F64 => "scatter_reduce_mean_div_f64", + _ => { + return Err(Error::UnsupportedDType { + dtype, + op: "scatter_reduce_mean_div", + }); + } + }; + + let func = get_kernel_function(&module, func_name)?; + + let grid = elementwise_launch_config(n); + let block = (BLOCK_SIZE, 1, 1); + let cfg = launch_config(grid, block, 0); + + let n_u32 = n as u32; + + let mut builder = stream.launch_builder(&func); + builder.arg(&sum_ptr); + builder.arg(&count_ptr); + builder.arg(&output_ptr); + builder.arg(&n_u32); + + builder.launch(cfg).map_err(|e| { + Error::Internal(format!( + "CUDA scatter_reduce_mean_div kernel launch failed: {:?}", + e + )) + })?; + + Ok(()) + } +} + // ============================================================================ // Gather 2D // ============================================================================ diff --git a/src/runtime/wgpu/ops/helpers.rs b/src/runtime/wgpu/ops/helpers.rs index 66b23df5..b70dd91a 100644 --- a/src/runtime/wgpu/ops/helpers.rs +++ b/src/runtime/wgpu/ops/helpers.rs @@ -701,6 +701,16 @@ pub(crate) struct ScatterReduceParams { pub(crate) _pad2: u32, } +/// Params for scatter_reduce mean division +#[repr(C)] +#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)] +pub(crate) struct MeanDivParams { + pub(crate) n: u32, + pub(crate) _pad0: u32, + pub(crate) _pad1: u32, + pub(crate) _pad2: u32, +} + /// Params for gather_2d operation #[repr(C)] #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)] diff --git a/src/runtime/wgpu/shaders/generator/index.rs b/src/runtime/wgpu/shaders/generator/index.rs index a3858ec7..9236c6c1 100644 --- a/src/runtime/wgpu/shaders/generator/index.rs +++ b/src/runtime/wgpu/shaders/generator/index.rs @@ -707,6 +707,234 @@ fn scatter_reduce_{op}_{suffix}(@builtin(global_invocation_id) gid: vec3) { } } +/// Generate WGSL shader for scatter_reduce prod operation. +/// +/// Uses CAS loop for atomic multiply (no native atomicMul in WGSL). +/// Only supports F32 (uses bitcast to u32 for atomics). +pub fn generate_scatter_reduce_prod_shader(dtype: DType) -> Result { + let t = wgsl_type(dtype)?; + let suffix = dtype_suffix(dtype)?; + let is_float = matches!(dtype, DType::F32); + + if is_float { + Ok(format!( + r#"// Auto-generated scatter_reduce_prod for {t} + +const WORKGROUP_SIZE: u32 = 256u; + +struct ScatterReduceParams {{ + dim: u32, + outer_size: u32, + dim_size: u32, + inner_size: u32, + src_dim_size: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +}} + +@group(0) @binding(0) var scatter_src: array<{t}>; +@group(0) @binding(1) var scatter_indices: array; +@group(0) @binding(2) var scatter_dst: array>; +@group(0) @binding(3) var scatter_params: ScatterReduceParams; + +@compute @workgroup_size(256) +fn scatter_reduce_prod_{suffix}(@builtin(global_invocation_id) gid: vec3) {{ + let idx = gid.x; + let total = scatter_params.outer_size * scatter_params.src_dim_size * scatter_params.inner_size; + if (idx >= total) {{ + return; + }} + + let inner = idx % scatter_params.inner_size; + let src_dim_idx = (idx / scatter_params.inner_size) % scatter_params.src_dim_size; + let outer = idx / (scatter_params.src_dim_size * scatter_params.inner_size); + + let index_val = scatter_indices[src_dim_idx]; + if (index_val < 0 || u32(index_val) >= scatter_params.dim_size) {{ + return; + }} + + let src_val = scatter_src[idx]; + let dst_idx = outer * scatter_params.dim_size * scatter_params.inner_size + u32(index_val) * scatter_params.inner_size + inner; + + // CAS loop for atomic multiply + var old_bits: u32; + var new_bits: u32; + loop {{ + old_bits = atomicLoad(&scatter_dst[dst_idx]); + let old_val = bitcast(old_bits); + let new_val = old_val * src_val; + new_bits = bitcast(new_val); + let result = atomicCompareExchangeWeak(&scatter_dst[dst_idx], old_bits, new_bits); + if (result.exchanged) {{ + break; + }} + }} +}} +"#, + t = t, + suffix = suffix, + )) + } else { + // Integer prod using CAS loop + let atomic_t = if dtype == DType::I32 { "i32" } else { "u32" }; + Ok(format!( + r#"// Auto-generated scatter_reduce_prod for {t} + +const WORKGROUP_SIZE: u32 = 256u; + +struct ScatterReduceParams {{ + dim: u32, + outer_size: u32, + dim_size: u32, + inner_size: u32, + src_dim_size: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +}} + +@group(0) @binding(0) var scatter_src: array<{t}>; +@group(0) @binding(1) var scatter_indices: array; +@group(0) @binding(2) var scatter_dst: array>; +@group(0) @binding(3) var scatter_params: ScatterReduceParams; + +@compute @workgroup_size(256) +fn scatter_reduce_prod_{suffix}(@builtin(global_invocation_id) gid: vec3) {{ + let idx = gid.x; + let total = scatter_params.outer_size * scatter_params.src_dim_size * scatter_params.inner_size; + if (idx >= total) {{ + return; + }} + + let inner = idx % scatter_params.inner_size; + let src_dim_idx = (idx / scatter_params.inner_size) % scatter_params.src_dim_size; + let outer = idx / (scatter_params.src_dim_size * scatter_params.inner_size); + + let index_val = scatter_indices[src_dim_idx]; + if (index_val < 0 || u32(index_val) >= scatter_params.dim_size) {{ + return; + }} + + let src_val = scatter_src[idx]; + let dst_idx = outer * scatter_params.dim_size * scatter_params.inner_size + u32(index_val) * scatter_params.inner_size + inner; + + // CAS loop for atomic multiply + loop {{ + let old_val = atomicLoad(&scatter_dst[dst_idx]); + let new_val = old_val * src_val; + let result = atomicCompareExchangeWeak(&scatter_dst[dst_idx], old_val, new_val); + if (result.exchanged) {{ + break; + }} + }} +}} +"#, + t = t, + suffix = suffix, + atomic_t = atomic_t, + )) + } +} + +/// Generate WGSL shader for scatter_reduce count (for mean computation). +/// +/// Atomically increments count buffer at scattered positions. +pub fn generate_scatter_reduce_count_shader(dtype: DType) -> Result { + let suffix = dtype_suffix(dtype)?; + + // Count buffer is always u32 (atomic) + Ok(format!( + r#"// Auto-generated scatter_reduce_count for mean computation + +const WORKGROUP_SIZE: u32 = 256u; + +struct ScatterReduceParams {{ + dim: u32, + outer_size: u32, + dim_size: u32, + inner_size: u32, + src_dim_size: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +}} + +@group(0) @binding(0) var scatter_indices: array; +@group(0) @binding(1) var scatter_count: array>; +@group(0) @binding(2) var scatter_params: ScatterReduceParams; + +@compute @workgroup_size(256) +fn scatter_reduce_count_{suffix}(@builtin(global_invocation_id) gid: vec3) {{ + let idx = gid.x; + let total = scatter_params.outer_size * scatter_params.src_dim_size * scatter_params.inner_size; + if (idx >= total) {{ + return; + }} + + let inner = idx % scatter_params.inner_size; + let src_dim_idx = (idx / scatter_params.inner_size) % scatter_params.src_dim_size; + let outer = idx / (scatter_params.src_dim_size * scatter_params.inner_size); + + let index_val = scatter_indices[src_dim_idx]; + if (index_val < 0 || u32(index_val) >= scatter_params.dim_size) {{ + return; + }} + + let dst_idx = outer * scatter_params.dim_size * scatter_params.inner_size + u32(index_val) * scatter_params.inner_size + inner; + + atomicAdd(&scatter_count[dst_idx], 1u); +}} +"#, + suffix = suffix, + )) +} + +/// Generate WGSL shader for scatter_reduce mean divide. +/// +/// Element-wise: output[i] = sum[i] / f32(count[i]). If count == 0, output = 0. +pub fn generate_scatter_reduce_mean_div_shader(dtype: DType) -> Result { + let t = wgsl_type(dtype)?; + let suffix = dtype_suffix(dtype)?; + + Ok(format!( + r#"// Auto-generated scatter_reduce_mean_div for {t} + +const WORKGROUP_SIZE: u32 = 256u; + +struct MeanDivParams {{ + n: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +}} + +@group(0) @binding(0) var mean_sum: array<{t}>; +@group(0) @binding(1) var mean_count: array; +@group(0) @binding(2) var mean_output: array<{t}>; +@group(0) @binding(3) var mean_params: MeanDivParams; + +@compute @workgroup_size(256) +fn scatter_reduce_mean_div_{suffix}(@builtin(global_invocation_id) gid: vec3) {{ + let idx = gid.x; + if (idx >= mean_params.n) {{ + return; + }} + + let c = mean_count[idx]; + if (c > 0u) {{ + mean_output[idx] = mean_sum[idx] / {t}(c); + }} else {{ + mean_output[idx] = {t}(0); + }} +}} +"#, + t = t, + suffix = suffix, + )) +} + /// Generate WGSL shader for index bounds validation. /// /// Validates that all indices are within bounds `[0, dim_size)`. diff --git a/src/runtime/wgpu/shaders/generator/mod.rs b/src/runtime/wgpu/shaders/generator/mod.rs index 8297aef1..36b43740 100644 --- a/src/runtime/wgpu/shaders/generator/mod.rs +++ b/src/runtime/wgpu/shaders/generator/mod.rs @@ -96,8 +96,9 @@ pub use fft::{ pub use index::{ generate_bincount_shader, generate_embedding_lookup_shader, generate_gather_2d_shader, generate_gather_nd_shader, generate_gather_shader, generate_index_put_shader, - generate_index_select_shader, generate_scatter_reduce_shader, generate_scatter_shader, - generate_validate_indices_shader, + generate_index_select_shader, generate_scatter_reduce_count_shader, + generate_scatter_reduce_mean_div_shader, generate_scatter_reduce_prod_shader, + generate_scatter_reduce_shader, generate_scatter_shader, generate_validate_indices_shader, }; pub use masked::{generate_masked_fill_shader, generate_masked_select_shader}; pub use matmul::{generate_matmul_bias_shader, generate_matmul_shader}; diff --git a/src/runtime/wgpu/shaders/index.rs b/src/runtime/wgpu/shaders/index.rs index 5236f4f0..66c9f9b5 100644 --- a/src/runtime/wgpu/shaders/index.rs +++ b/src/runtime/wgpu/shaders/index.rs @@ -78,6 +78,11 @@ fn kernel_name(op: &'static str, dtype: DType) -> Result<&'static str> { ("scatter_reduce_min", DType::F32) => Ok("scatter_reduce_min_f32"), ("scatter_reduce_min", DType::I32) => Ok("scatter_reduce_min_i32"), ("scatter_reduce_min", DType::U32) => Ok("scatter_reduce_min_u32"), + ("scatter_reduce_prod", DType::F32) => Ok("scatter_reduce_prod_f32"), + ("scatter_reduce_prod", DType::I32) => Ok("scatter_reduce_prod_i32"), + ("scatter_reduce_prod", DType::U32) => Ok("scatter_reduce_prod_u32"), + ("scatter_reduce_count", DType::F32) => Ok("scatter_reduce_count_f32"), + ("scatter_reduce_mean_div", DType::F32) => Ok("scatter_reduce_mean_div_f32"), ("gather_2d", DType::F32) => Ok("gather_2d_f32"), ("gather_2d", DType::I32) => Ok("gather_2d_i32"), ("gather_2d", DType::U32) => Ok("gather_2d_u32"), @@ -161,8 +166,8 @@ pub fn launch_index_put( let shader_source = generate_index_put_shader(dtype)?; let module = cache.get_or_create_module(name, &shader_source); let layout = cache.get_or_create_layout(LayoutKey { - num_storage_buffers: 4, - num_uniform_buffers: 0, + num_storage_buffers: 3, + num_uniform_buffers: 1, num_readonly_storage: 0, }); let pipeline = cache.get_or_create_pipeline(name, name, &module, &layout); @@ -785,6 +790,154 @@ pub fn launch_scatter_reduce( Ok(()) } +// ============================================================================ +// Scatter Reduce Prod Operation +// ============================================================================ + +/// Launch a scatter_reduce_prod operation kernel. +/// +/// Scatters values with product reduction using CAS loop. +pub fn launch_scatter_reduce_prod( + cache: &PipelineCache, + queue: &Queue, + src: &Buffer, + indices: &Buffer, + dst: &Buffer, + params_buffer: &Buffer, + total_src: usize, + dtype: DType, +) -> Result<()> { + check_dtype_supported(dtype, "scatter_reduce_prod")?; + + let name = kernel_name("scatter_reduce_prod", dtype)?; + let shader_source = super::generator::generate_scatter_reduce_prod_shader(dtype)?; + let module = cache.get_or_create_module(name, &shader_source); + let layout = cache.get_or_create_layout(LayoutKey { + num_storage_buffers: 3, + num_uniform_buffers: 1, + num_readonly_storage: 0, + }); + let pipeline = cache.get_or_create_pipeline(name, name, &module, &layout); + + let bind_group = cache.create_bind_group(&layout, &[src, indices, dst, params_buffer]); + + let mut encoder = cache + .device() + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("scatter_reduce_prod"), + }); + + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("scatter_reduce_prod"), + timestamp_writes: None, + }); + pass.set_pipeline(&pipeline); + pass.set_bind_group(0, Some(&bind_group), &[]); + pass.dispatch_workgroups(workgroup_count(total_src), 1, 1); + } + + queue.submit(std::iter::once(encoder.finish())); + Ok(()) +} + +// ============================================================================ +// Scatter Reduce Count Operation (for Mean) +// ============================================================================ + +/// Launch a scatter_reduce_count operation kernel. +/// +/// Atomically counts scattered elements per destination position. +pub fn launch_scatter_reduce_count( + cache: &PipelineCache, + queue: &Queue, + indices: &Buffer, + count: &Buffer, + params_buffer: &Buffer, + total_src: usize, + dtype: DType, +) -> Result<()> { + let name = kernel_name("scatter_reduce_count", dtype)?; + let shader_source = super::generator::generate_scatter_reduce_count_shader(dtype)?; + let module = cache.get_or_create_module(name, &shader_source); + let layout = cache.get_or_create_layout(LayoutKey { + num_storage_buffers: 2, + num_uniform_buffers: 1, + num_readonly_storage: 0, + }); + let pipeline = cache.get_or_create_pipeline(name, name, &module, &layout); + + let bind_group = cache.create_bind_group(&layout, &[indices, count, params_buffer]); + + let mut encoder = cache + .device() + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("scatter_reduce_count"), + }); + + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("scatter_reduce_count"), + timestamp_writes: None, + }); + pass.set_pipeline(&pipeline); + pass.set_bind_group(0, Some(&bind_group), &[]); + pass.dispatch_workgroups(workgroup_count(total_src), 1, 1); + } + + queue.submit(std::iter::once(encoder.finish())); + Ok(()) +} + +// ============================================================================ +// Scatter Reduce Mean Divide Operation +// ============================================================================ + +/// Launch scatter_reduce_mean_div: output[i] = sum[i] / count[i]. +pub fn launch_scatter_reduce_mean_div( + cache: &PipelineCache, + queue: &Queue, + sum_buf: &Buffer, + count_buf: &Buffer, + output: &Buffer, + params_buffer: &Buffer, + n: usize, + dtype: DType, +) -> Result<()> { + check_dtype_supported(dtype, "scatter_reduce_mean_div")?; + + let name = kernel_name("scatter_reduce_mean_div", dtype)?; + let shader_source = super::generator::generate_scatter_reduce_mean_div_shader(dtype)?; + let module = cache.get_or_create_module(name, &shader_source); + let layout = cache.get_or_create_layout(LayoutKey { + num_storage_buffers: 3, + num_uniform_buffers: 1, + num_readonly_storage: 0, + }); + let pipeline = cache.get_or_create_pipeline(name, name, &module, &layout); + + let bind_group = cache.create_bind_group(&layout, &[sum_buf, count_buf, output, params_buffer]); + + let mut encoder = cache + .device() + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("scatter_reduce_mean_div"), + }); + + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("scatter_reduce_mean_div"), + timestamp_writes: None, + }); + pass.set_pipeline(&pipeline); + pass.set_bind_group(0, Some(&bind_group), &[]); + pass.dispatch_workgroups(workgroup_count(n), 1, 1); + } + + queue.submit(std::iter::once(encoder.finish())); + Ok(()) +} + // ============================================================================ // Embedding Lookup Operation // ============================================================================ @@ -873,7 +1026,7 @@ pub fn launch_gather_2d( let layout = cache.get_or_create_layout(LayoutKey { num_storage_buffers: 4, num_uniform_buffers: 1, - num_readonly_storage: 0, + num_readonly_storage: 3, }); let pipeline = cache.get_or_create_pipeline(name, name, &module, &layout); diff --git a/src/runtime/wgpu/shaders/mod.rs b/src/runtime/wgpu/shaders/mod.rs index 59316a72..f96ea7f4 100644 --- a/src/runtime/wgpu/shaders/mod.rs +++ b/src/runtime/wgpu/shaders/mod.rs @@ -117,7 +117,10 @@ pub use generator::{ }; #[cfg(feature = "sparse")] pub use generator::{generate_csr_spmm_shader, generate_csr_spmv_shader}; -pub use index::{launch_bincount, launch_gather_2d, launch_gather_nd, launch_scatter_reduce}; +pub use index::{ + launch_bincount, launch_gather_2d, launch_gather_nd, launch_scatter_reduce, + launch_scatter_reduce_count, launch_scatter_reduce_mean_div, launch_scatter_reduce_prod, +}; pub use logical::{launch_logical_and, launch_logical_not, launch_logical_or, launch_logical_xor}; pub use matrix_funcs_launcher::{ compute_schur_func_gpu, launch_diagonal_func, launch_parlett_column, diff --git a/tests/backend_parity/complex.rs b/tests/backend_parity/complex.rs index 221ba570..3f699184 100644 --- a/tests/backend_parity/complex.rs +++ b/tests/backend_parity/complex.rs @@ -11,7 +11,7 @@ use numr::ops::{BinaryOps, ComplexOps, UnaryOps}; #[cfg(feature = "wgpu")] use numr::prelude::DType; use numr::runtime::Runtime; -use numr::runtime::cpu::{CpuDevice, CpuRuntime}; +use numr::runtime::cpu::{CpuClient, CpuDevice, CpuRuntime, ParallelismConfig}; use numr::tensor::Tensor; fn assert_complex_close(cpu: &[Complex64], other: &[Complex64], tol: f32, label: &str) { @@ -38,6 +38,100 @@ fn assert_complex_close(cpu: &[Complex64], other: &[Complex64], tol: f32, label: } } +#[test] +fn test_cpu_complex_parallelism_config_matches_default() { + let device = CpuDevice::new(); + let default_client = CpuClient::new(device.clone()); + let configured_client = + default_client.with_parallelism(ParallelismConfig::new(Some(2), Some(257))); + + // Keep this above CPU complex kernel parallel threshold (4096) to exercise + // the Rayon chunking path with custom chunk_size. + let shape = [128, 64]; + let numel: usize = shape.iter().product(); + let a_data: Vec = (0..numel) + .map(|i| Complex64::new((i as f32 * 0.011).sin(), (i as f32 * 0.017).cos())) + .collect(); + let real_data: Vec = (0..numel) + .map(|i| 1.1 + (i as f32 * 0.019).sin().abs()) + .collect(); + let imag_data: Vec = (0..numel).map(|i| (i as f32 * 0.023).cos()).collect(); + + let a = Tensor::::from_slice(&a_data, &shape, &device); + let real = Tensor::::from_slice(&real_data, &shape, &device); + let imag = Tensor::::from_slice(&imag_data, &shape, &device); + + let base_conj: Vec = default_client.conj(&a).unwrap().to_vec(); + let cfg_conj: Vec = configured_client.conj(&a).unwrap().to_vec(); + assert_complex_close( + &base_conj, + &cfg_conj, + 1e-6, + "cpu complex conj parallelism config", + ); + + let base_real: Vec = default_client.real(&a).unwrap().to_vec(); + let cfg_real: Vec = configured_client.real(&a).unwrap().to_vec(); + for (idx, (b, c)) in base_real.iter().zip(cfg_real.iter()).enumerate() { + assert!((b - c).abs() <= 1e-6, "cpu complex real mismatch at {idx}"); + } + + let base_imag: Vec = default_client.imag(&a).unwrap().to_vec(); + let cfg_imag: Vec = configured_client.imag(&a).unwrap().to_vec(); + for (idx, (b, c)) in base_imag.iter().zip(cfg_imag.iter()).enumerate() { + assert!((b - c).abs() <= 1e-6, "cpu complex imag mismatch at {idx}"); + } + + let base_angle: Vec = default_client.angle(&a).unwrap().to_vec(); + let cfg_angle: Vec = configured_client.angle(&a).unwrap().to_vec(); + for (idx, (b, c)) in base_angle.iter().zip(cfg_angle.iter()).enumerate() { + assert!((b - c).abs() <= 1e-6, "cpu complex angle mismatch at {idx}"); + } + + let base_make: Vec = default_client.make_complex(&real, &imag).unwrap().to_vec(); + let cfg_make: Vec = configured_client + .make_complex(&real, &imag) + .unwrap() + .to_vec(); + assert_complex_close( + &base_make, + &cfg_make, + 1e-6, + "cpu make_complex parallelism config", + ); + + let made = default_client.make_complex(&real, &imag).unwrap(); + let base_mul: Vec = default_client + .complex_mul_real(&made, &real) + .unwrap() + .to_vec(); + let cfg_mul: Vec = configured_client + .complex_mul_real(&made, &real) + .unwrap() + .to_vec(); + assert_complex_close( + &base_mul, + &cfg_mul, + 1e-6, + "cpu complex_mul_real parallelism config", + ); + + let base_div: Vec = default_client + .complex_div_real(&made, &real) + .unwrap() + .to_vec(); + let cfg_div: Vec = configured_client + .complex_div_real(&made, &real) + .unwrap() + .to_vec(); + assert_complex_close( + &base_div, + &cfg_div, + 1e-6, + "cpu complex_div_real parallelism config", + ); +} + #[test] fn test_complex_angle_parity() { let cpu_device = CpuDevice::new(); diff --git a/tests/backend_parity/fft.rs b/tests/backend_parity/fft.rs index 8cb2ce8e..6be1e4da 100644 --- a/tests/backend_parity/fft.rs +++ b/tests/backend_parity/fft.rs @@ -7,7 +7,7 @@ use crate::backend_parity::helpers::with_wgpu_backend; use numr::algorithm::fft::{FftAlgorithms, FftDirection, FftNormalization}; use numr::dtype::Complex64; use numr::runtime::RuntimeClient; -use numr::runtime::cpu::{CpuClient, CpuDevice, CpuRuntime}; +use numr::runtime::cpu::{CpuClient, CpuDevice, CpuRuntime, ParallelismConfig}; use numr::tensor::Tensor; fn get_cpu_client() -> CpuClient { @@ -23,6 +23,13 @@ fn assert_complex_close(cpu: &[Complex64], other: &[Complex64], tol: f32, label: } } +fn assert_f32_close(cpu: &[f32], other: &[f32], tol: f32, label: &str) { + assert_eq!(cpu.len(), other.len(), "{} length mismatch", label); + for (i, (c, g)) in cpu.iter().zip(other.iter()).enumerate() { + assert!((c - g).abs() < tol, "{} idx {}", label, i); + } +} + #[test] fn test_fft_forward_parity() { let cpu_client = get_cpu_client(); @@ -201,3 +208,113 @@ fn test_fftshift_parity() { assert_complex_close(&cpu_data, &data, 1e-5, "fftshift wgpu"); }); } + +#[test] +fn test_cpu_fft_parallelism_config_matches_default() { + let device = CpuDevice::new(); + let default_client = CpuClient::new(device.clone()); + let configured_client = + default_client.with_parallelism(ParallelismConfig::new(Some(1), Some(1024))); + + // Use a batched shape so CPU FFT goes through the batched kernel path. + let shape = [8, 128]; + let numel: usize = shape.iter().product(); + let input_data: Vec = (0..numel) + .map(|i| Complex64::new((i as f32 * 0.031).sin(), (i as f32 * 0.017).cos())) + .collect(); + + let input = Tensor::::from_slice(&input_data, &shape, &device); + let base = default_client + .fft(&input, FftDirection::Forward, FftNormalization::None) + .unwrap(); + let configured = configured_client + .fft(&input, FftDirection::Forward, FftNormalization::None) + .unwrap(); + + let base_data: Vec = base.to_vec(); + let configured_data: Vec = configured.to_vec(); + assert_complex_close( + &base_data, + &configured_data, + 1e-5, + "cpu fft parallelism config", + ); +} + +#[test] +fn test_cpu_rfft_irfft_parallelism_config_matches_default() { + let device = CpuDevice::new(); + let default_client = CpuClient::new(device.clone()); + let configured_client = + default_client.with_parallelism(ParallelismConfig::new(Some(1), Some(1024))); + + let shape = [6, 64]; + let numel: usize = shape.iter().product(); + let input_data: Vec = (0..numel).map(|i| (i as f32 * 0.023).sin()).collect(); + + let input = Tensor::::from_slice(&input_data, &shape, &device); + let base_freq = default_client.rfft(&input, FftNormalization::None).unwrap(); + let cfg_freq = configured_client + .rfft(&input, FftNormalization::None) + .unwrap(); + let base_freq_data: Vec = base_freq.to_vec(); + let cfg_freq_data: Vec = cfg_freq.to_vec(); + assert_complex_close( + &base_freq_data, + &cfg_freq_data, + 1e-5, + "cpu rfft parallelism config", + ); + + let base_rec = default_client + .irfft(&base_freq, Some(shape[1]), FftNormalization::Backward) + .unwrap(); + let cfg_rec = configured_client + .irfft(&cfg_freq, Some(shape[1]), FftNormalization::Backward) + .unwrap(); + let base_rec_data: Vec = base_rec.to_vec(); + let cfg_rec_data: Vec = cfg_rec.to_vec(); + assert_f32_close( + &base_rec_data, + &cfg_rec_data, + 1e-5, + "cpu irfft parallelism config", + ); +} + +#[test] +fn test_cpu_fftshift_parallelism_config_matches_default() { + let device = CpuDevice::new(); + let default_client = CpuClient::new(device.clone()); + let configured_client = + default_client.with_parallelism(ParallelismConfig::new(Some(1), Some(1024))); + + let shape = [5, 32]; + let numel: usize = shape.iter().product(); + let input_data: Vec = (0..numel) + .map(|i| Complex64::new((i as f32 * 0.013).cos(), (i as f32 * 0.019).sin())) + .collect(); + + let input = Tensor::::from_slice(&input_data, &shape, &device); + let base_shift = default_client.fftshift(&input).unwrap(); + let cfg_shift = configured_client.fftshift(&input).unwrap(); + let base_shift_data: Vec = base_shift.to_vec(); + let cfg_shift_data: Vec = cfg_shift.to_vec(); + assert_complex_close( + &base_shift_data, + &cfg_shift_data, + 1e-5, + "cpu fftshift parallelism config", + ); + + let base_unshift = default_client.ifftshift(&base_shift).unwrap(); + let cfg_unshift = configured_client.ifftshift(&cfg_shift).unwrap(); + let base_unshift_data: Vec = base_unshift.to_vec(); + let cfg_unshift_data: Vec = cfg_unshift.to_vec(); + assert_complex_close( + &base_unshift_data, + &cfg_unshift_data, + 1e-5, + "cpu ifftshift parallelism config", + ); +} diff --git a/tests/backend_parity/indexing.rs b/tests/backend_parity/indexing.rs index 7258fb08..c33f4340 100644 --- a/tests/backend_parity/indexing.rs +++ b/tests/backend_parity/indexing.rs @@ -4,6 +4,8 @@ use crate::backend_parity::helpers::with_cuda_backend; #[cfg(feature = "wgpu")] use crate::backend_parity::helpers::with_wgpu_backend; +use crate::common::create_cpu_client; +use numr::error::Error; use numr::ops::IndexingOps; #[cfg(feature = "cuda")] use numr::runtime::Runtime; @@ -117,3 +119,173 @@ fn test_masked_ops_parity() { assert_eq!(filled, vec![-1.0, 2.0, -1.0, 4.0, 5.0, -1.0, 7.0, -1.0]); }); } + +#[test] +fn test_take_put_parity() { + let (cpu_client, cpu_device) = create_cpu_client(); + let a_cpu = Tensor::from_slice( + &[10.0f32, 20.0, 30.0, 40.0, 50.0, 60.0], + &[2, 3], + &cpu_device, + ); + let idx_cpu = Tensor::from_slice(&[5i32, 0, 2, 4], &[2, 2], &cpu_device); + let put_values_cpu = Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2], &cpu_device); + let cpu_take: Vec = cpu_client.take(&a_cpu, &idx_cpu).unwrap().to_vec(); + let cpu_put: Vec = cpu_client + .put(&a_cpu, &idx_cpu, &put_values_cpu) + .unwrap() + .to_vec(); + assert_eq!(cpu_take, vec![60.0, 10.0, 30.0, 50.0]); + assert_eq!(cpu_put, vec![2.0, 20.0, 3.0, 40.0, 4.0, 1.0]); + + #[cfg(feature = "cuda")] + with_cuda_backend(|cuda_client, cuda_device| { + let a = Tensor::::from_slice( + &[10.0f32, 20.0, 30.0, 40.0, 50.0, 60.0], + &[2, 3], + &cuda_device, + ); + let idx = Tensor::::from_slice( + &[5i32, 0, 2, 4], + &[2, 2], + &cuda_device, + ); + let put_values = Tensor::::from_slice( + &[1.0f32, 2.0, 3.0, 4.0], + &[2, 2], + &cuda_device, + ); + + let take: Vec = cuda_client.take(&a, &idx).unwrap().to_vec(); + assert_eq!(cpu_take, take); + + let put: Vec = cuda_client.put(&a, &idx, &put_values).unwrap().to_vec(); + assert_eq!(cpu_put, put); + }); + + #[cfg(feature = "wgpu")] + with_wgpu_backend(|wgpu_client, wgpu_device| { + let a = Tensor::::from_slice( + &[10.0f32, 20.0, 30.0, 40.0, 50.0, 60.0], + &[2, 3], + &wgpu_device, + ); + let idx = Tensor::::from_slice( + &[5i32, 0, 2, 4], + &[2, 2], + &wgpu_device, + ); + let put_values = Tensor::::from_slice( + &[1.0f32, 2.0, 3.0, 4.0], + &[2, 2], + &wgpu_device, + ); + + let take: Vec = wgpu_client.take(&a, &idx).unwrap().to_vec(); + assert_eq!(take, vec![60.0, 10.0, 30.0, 50.0]); + + let put: Vec = wgpu_client.put(&a, &idx, &put_values).unwrap().to_vec(); + assert_eq!(put, vec![2.0, 20.0, 3.0, 40.0, 4.0, 1.0]); + }); +} + +#[test] +fn test_take_put_i64_indices_parity() { + let (cpu_client, cpu_device) = create_cpu_client(); + let a_cpu = Tensor::from_slice( + &[10.0f32, 20.0, 30.0, 40.0, 50.0, 60.0], + &[2, 3], + &cpu_device, + ); + let idx_cpu = Tensor::from_slice(&[5i64, 0, 2, 4], &[2, 2], &cpu_device); + let put_values_cpu = Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2], &cpu_device); + let cpu_take: Vec = cpu_client.take(&a_cpu, &idx_cpu).unwrap().to_vec(); + let cpu_put: Vec = cpu_client + .put(&a_cpu, &idx_cpu, &put_values_cpu) + .unwrap() + .to_vec(); + assert_eq!(cpu_take, vec![60.0, 10.0, 30.0, 50.0]); + assert_eq!(cpu_put, vec![2.0, 20.0, 3.0, 40.0, 4.0, 1.0]); + + #[cfg(feature = "cuda")] + with_cuda_backend(|cuda_client, cuda_device| { + let a = Tensor::::from_slice( + &[10.0f32, 20.0, 30.0, 40.0, 50.0, 60.0], + &[2, 3], + &cuda_device, + ); + let idx = Tensor::::from_slice( + &[5i64, 0, 2, 4], + &[2, 2], + &cuda_device, + ); + let put_values = Tensor::::from_slice( + &[1.0f32, 2.0, 3.0, 4.0], + &[2, 2], + &cuda_device, + ); + + let take: Vec = cuda_client.take(&a, &idx).unwrap().to_vec(); + assert_eq!(cpu_take, take); + + let put: Vec = cuda_client.put(&a, &idx, &put_values).unwrap().to_vec(); + assert_eq!(cpu_put, put); + }); + + #[cfg(feature = "wgpu")] + with_wgpu_backend(|wgpu_client, wgpu_device| { + let a = Tensor::::from_slice( + &[10.0f32, 20.0, 30.0, 40.0, 50.0, 60.0], + &[2, 3], + &wgpu_device, + ); + let idx = Tensor::::from_slice( + &[5i64, 0, 2, 4], + &[2, 2], + &wgpu_device, + ); + let put_values = Tensor::::from_slice( + &[1.0f32, 2.0, 3.0, 4.0], + &[2, 2], + &wgpu_device, + ); + + let take: Vec = wgpu_client.take(&a, &idx).unwrap().to_vec(); + assert_eq!(take, vec![60.0, 10.0, 30.0, 50.0]); + + let put: Vec = wgpu_client.put(&a, &idx, &put_values).unwrap().to_vec(); + assert_eq!(put, vec![2.0, 20.0, 3.0, 40.0, 4.0, 1.0]); + }); +} + +#[test] +fn test_take_put_reject_non_integer_indices() { + let (cpu_client, cpu_device) = create_cpu_client(); + let a_cpu = Tensor::from_slice( + &[10.0f32, 20.0, 30.0, 40.0, 50.0, 60.0], + &[2, 3], + &cpu_device, + ); + let idx_cpu = Tensor::from_slice(&[0.0f32, 2.0], &[2], &cpu_device); + let put_values_cpu = Tensor::from_slice(&[1.0f32, 2.0], &[2], &cpu_device); + + let take_err = cpu_client.take(&a_cpu, &idx_cpu).unwrap_err(); + match take_err { + Error::InvalidArgument { arg, reason } => { + assert_eq!(arg, "indices"); + assert!(reason.contains("I32 or I64")); + } + other => panic!("unexpected error variant: {other:?}"), + } + + let put_err = cpu_client + .put(&a_cpu, &idx_cpu, &put_values_cpu) + .unwrap_err(); + match put_err { + Error::InvalidArgument { arg, reason } => { + assert_eq!(arg, "indices"); + assert!(reason.contains("I32 or I64")); + } + other => panic!("unexpected error variant: {other:?}"), + } +} diff --git a/tests/backend_parity/indexing_advanced.rs b/tests/backend_parity/indexing_advanced.rs index f17ee6b2..ac59a948 100644 --- a/tests/backend_parity/indexing_advanced.rs +++ b/tests/backend_parity/indexing_advanced.rs @@ -37,6 +37,254 @@ fn test_index_select_parity() { }); } +#[test] +fn test_i32_indices_parity() { + let (cpu_client, cpu_device) = create_cpu_client(); + + let input = Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2], &cpu_device); + let idx_1d = Tensor::from_slice(&[2i32, 0], &[2], &cpu_device); + let idx_2d = Tensor::from_slice(&[0i32, 2, 1, 0], &[2, 2], &cpu_device); + + let cpu_index_select: Vec = cpu_client + .index_select(&input, 0, &idx_1d) + .unwrap() + .to_vec(); + let cpu_gather: Vec = cpu_client.gather(&input, 0, &idx_2d).unwrap().to_vec(); + let cpu_scatter: Vec = cpu_client + .scatter( + &Tensor::from_slice(&[0.0f32; 6], &[3, 2], &cpu_device), + 0, + &idx_2d, + &Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2], &cpu_device), + ) + .unwrap() + .to_vec(); + let cpu_index_put: Vec = cpu_client + .index_put( + &input, + 0, + &idx_1d, + &Tensor::from_slice(&[10.0f32, 11.0, 12.0, 13.0], &[2, 2], &cpu_device), + ) + .unwrap() + .to_vec(); + + let nd_input = Tensor::from_slice(&[0.0f32, 1.0, 2.0, 3.0], &[2, 2], &cpu_device); + let nd_idx = Tensor::from_slice(&[0i32, 0, 1, 1], &[2, 2], &cpu_device); + let cpu_gather_nd: Vec = cpu_client.gather_nd(&nd_input, &nd_idx).unwrap().to_vec(); + + let emb = Tensor::from_slice( + &[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], + &[4, 2], + &cpu_device, + ); + let emb_idx = Tensor::from_slice(&[3i32, 0, 1], &[3], &cpu_device); + let cpu_emb: Vec = cpu_client + .embedding_lookup(&emb, &emb_idx) + .unwrap() + .to_vec(); + + let g2d_input = Tensor::from_slice( + &[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], + &[3, 3], + &cpu_device, + ); + let g2d_rows = Tensor::from_slice(&[0i32, 1, 2, 0], &[4], &cpu_device); + let g2d_cols = Tensor::from_slice(&[0i32, 1, 2, 2], &[4], &cpu_device); + let cpu_g2d: Vec = cpu_client + .gather_2d(&g2d_input, &g2d_rows, &g2d_cols) + .unwrap() + .to_vec(); + + let cpu_scatter_reduce: Vec = cpu_client + .scatter_reduce( + &Tensor::from_slice(&[0.0f32, 0.0, 0.0, 0.0], &[4], &cpu_device), + 0, + &Tensor::from_slice(&[0i32, 0, 2], &[3], &cpu_device), + &Tensor::from_slice(&[1.0f32, 2.0, 3.0], &[3], &cpu_device), + ScatterReduceOp::Sum, + false, + ) + .unwrap() + .to_vec(); + + #[cfg(feature = "cuda")] + with_cuda_backend(|cuda_client, cuda_device| { + let input = Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2], &cuda_device); + let idx_1d = Tensor::from_slice(&[2i32, 0], &[2], &cuda_device); + let idx_2d = Tensor::from_slice(&[0i32, 2, 1, 0], &[2, 2], &cuda_device); + + let got_index_select: Vec = cuda_client + .index_select(&input, 0, &idx_1d) + .unwrap() + .to_vec(); + assert_parity_f32( + &cpu_index_select, + &got_index_select, + "index_select_i32_cuda", + ); + + let got_gather: Vec = cuda_client.gather(&input, 0, &idx_2d).unwrap().to_vec(); + assert_parity_f32(&cpu_gather, &got_gather, "gather_i32_cuda"); + + let got_scatter: Vec = cuda_client + .scatter( + &Tensor::from_slice(&[0.0f32; 6], &[3, 2], &cuda_device), + 0, + &idx_2d, + &Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2], &cuda_device), + ) + .unwrap() + .to_vec(); + assert_parity_f32(&cpu_scatter, &got_scatter, "scatter_i32_cuda"); + let got_index_put: Vec = cuda_client + .index_put( + &input, + 0, + &idx_1d, + &Tensor::from_slice(&[10.0f32, 11.0, 12.0, 13.0], &[2, 2], &cuda_device), + ) + .unwrap() + .to_vec(); + assert_parity_f32(&cpu_index_put, &got_index_put, "index_put_i32_cuda"); + + let nd_input = Tensor::from_slice(&[0.0f32, 1.0, 2.0, 3.0], &[2, 2], &cuda_device); + let nd_idx = Tensor::from_slice(&[0i32, 0, 1, 1], &[2, 2], &cuda_device); + let got_gather_nd: Vec = cuda_client.gather_nd(&nd_input, &nd_idx).unwrap().to_vec(); + assert_parity_f32(&cpu_gather_nd, &got_gather_nd, "gather_nd_i32_cuda"); + + let emb = Tensor::from_slice( + &[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], + &[4, 2], + &cuda_device, + ); + let emb_idx = Tensor::from_slice(&[3i32, 0, 1], &[3], &cuda_device); + let got_emb: Vec = cuda_client + .embedding_lookup(&emb, &emb_idx) + .unwrap() + .to_vec(); + assert_parity_f32(&cpu_emb, &got_emb, "embedding_i32_cuda"); + + let g2d_input = Tensor::from_slice( + &[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], + &[3, 3], + &cuda_device, + ); + let g2d_rows = Tensor::from_slice(&[0i32, 1, 2, 0], &[4], &cuda_device); + let g2d_cols = Tensor::from_slice(&[0i32, 1, 2, 2], &[4], &cuda_device); + let got_g2d: Vec = cuda_client + .gather_2d(&g2d_input, &g2d_rows, &g2d_cols) + .unwrap() + .to_vec(); + assert_parity_f32(&cpu_g2d, &got_g2d, "gather_2d_i32_cuda"); + + let got_scatter_reduce: Vec = cuda_client + .scatter_reduce( + &Tensor::from_slice(&[0.0f32, 0.0, 0.0, 0.0], &[4], &cuda_device), + 0, + &Tensor::from_slice(&[0i32, 0, 2], &[3], &cuda_device), + &Tensor::from_slice(&[1.0f32, 2.0, 3.0], &[3], &cuda_device), + ScatterReduceOp::Sum, + false, + ) + .unwrap() + .to_vec(); + assert_parity_f32( + &cpu_scatter_reduce, + &got_scatter_reduce, + "scatter_reduce_i32_cuda", + ); + }); + + #[cfg(feature = "wgpu")] + with_wgpu_backend(|wgpu_client, wgpu_device| { + let input = Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2], &wgpu_device); + let idx_1d = Tensor::from_slice(&[2i32, 0], &[2], &wgpu_device); + let idx_2d = Tensor::from_slice(&[0i32, 2, 1, 0], &[2, 2], &wgpu_device); + + let got_index_select: Vec = wgpu_client + .index_select(&input, 0, &idx_1d) + .unwrap() + .to_vec(); + assert_parity_f32( + &cpu_index_select, + &got_index_select, + "index_select_i32_wgpu", + ); + + let got_gather: Vec = wgpu_client.gather(&input, 0, &idx_2d).unwrap().to_vec(); + assert_parity_f32(&cpu_gather, &got_gather, "gather_i32_wgpu"); + + let got_scatter: Vec = wgpu_client + .scatter( + &Tensor::from_slice(&[0.0f32; 6], &[3, 2], &wgpu_device), + 0, + &idx_2d, + &Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0], &[2, 2], &wgpu_device), + ) + .unwrap() + .to_vec(); + assert_parity_f32(&cpu_scatter, &got_scatter, "scatter_i32_wgpu"); + let got_index_put: Vec = wgpu_client + .index_put( + &input, + 0, + &idx_1d, + &Tensor::from_slice(&[10.0f32, 11.0, 12.0, 13.0], &[2, 2], &wgpu_device), + ) + .unwrap() + .to_vec(); + assert_parity_f32(&cpu_index_put, &got_index_put, "index_put_i32_wgpu"); + + let nd_input = Tensor::from_slice(&[0.0f32, 1.0, 2.0, 3.0], &[2, 2], &wgpu_device); + let nd_idx = Tensor::from_slice(&[0i32, 0, 1, 1], &[2, 2], &wgpu_device); + let got_gather_nd: Vec = wgpu_client.gather_nd(&nd_input, &nd_idx).unwrap().to_vec(); + assert_parity_f32(&cpu_gather_nd, &got_gather_nd, "gather_nd_i32_wgpu"); + + let emb = Tensor::from_slice( + &[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], + &[4, 2], + &wgpu_device, + ); + let emb_idx = Tensor::from_slice(&[3i32, 0, 1], &[3], &wgpu_device); + let got_emb: Vec = wgpu_client + .embedding_lookup(&emb, &emb_idx) + .unwrap() + .to_vec(); + assert_parity_f32(&cpu_emb, &got_emb, "embedding_i32_wgpu"); + + let g2d_input = Tensor::from_slice( + &[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], + &[3, 3], + &wgpu_device, + ); + let g2d_rows = Tensor::from_slice(&[0i32, 1, 2, 0], &[4], &wgpu_device); + let g2d_cols = Tensor::from_slice(&[0i32, 1, 2, 2], &[4], &wgpu_device); + let got_g2d: Vec = wgpu_client + .gather_2d(&g2d_input, &g2d_rows, &g2d_cols) + .unwrap() + .to_vec(); + assert_parity_f32(&cpu_g2d, &got_g2d, "gather_2d_i32_wgpu"); + + let got_scatter_reduce: Vec = wgpu_client + .scatter_reduce( + &Tensor::from_slice(&[0.0f32, 0.0, 0.0, 0.0], &[4], &wgpu_device), + 0, + &Tensor::from_slice(&[0i32, 0, 2], &[3], &wgpu_device), + &Tensor::from_slice(&[1.0f32, 2.0, 3.0], &[3], &wgpu_device), + ScatterReduceOp::Sum, + false, + ) + .unwrap() + .to_vec(); + assert_parity_f32( + &cpu_scatter_reduce, + &got_scatter_reduce, + "scatter_reduce_i32_wgpu", + ); + }); +} + #[test] fn test_gather_scatter_parity() { let input = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; @@ -182,3 +430,58 @@ fn test_scatter_reduce_sum_parity() { assert_parity_f32(&cpu, &got, "scatter_reduce_sum_wgpu"); }); } + +#[test] +fn test_scatter_reduce_mean_prod_parity() { + let (cpu_client, cpu_device) = create_cpu_client(); + let dst = Tensor::from_slice(&[10.0f32, 20.0, 30.0, 40.0], &[4], &cpu_device); + let idx = Tensor::from_slice(&[0i64, 0, 2], &[3], &cpu_device); + let src = Tensor::from_slice(&[2.0f32, 4.0, 8.0], &[3], &cpu_device); + + let cpu_mean: Vec = cpu_client + .scatter_reduce(&dst, 0, &idx, &src, ScatterReduceOp::Mean, true) + .unwrap() + .to_vec(); + let cpu_prod: Vec = cpu_client + .scatter_reduce(&dst, 0, &idx, &src, ScatterReduceOp::Prod, true) + .unwrap() + .to_vec(); + + #[cfg(feature = "cuda")] + with_cuda_backend(|cuda_client, cuda_device| { + let d = Tensor::from_slice(&[10.0f32, 20.0, 30.0, 40.0], &[4], &cuda_device); + let i = Tensor::from_slice(&[0i64, 0, 2], &[3], &cuda_device); + let s = Tensor::from_slice(&[2.0f32, 4.0, 8.0], &[3], &cuda_device); + + let mean: Vec = cuda_client + .scatter_reduce(&d, 0, &i, &s, ScatterReduceOp::Mean, true) + .unwrap() + .to_vec(); + assert_parity_f32(&cpu_mean, &mean, "scatter_reduce_mean_cuda"); + + let prod: Vec = cuda_client + .scatter_reduce(&d, 0, &i, &s, ScatterReduceOp::Prod, true) + .unwrap() + .to_vec(); + assert_parity_f32(&cpu_prod, &prod, "scatter_reduce_prod_cuda"); + }); + + #[cfg(feature = "wgpu")] + with_wgpu_backend(|wgpu_client, wgpu_device| { + let d = Tensor::from_slice(&[10.0f32, 20.0, 30.0, 40.0], &[4], &wgpu_device); + let i = Tensor::from_slice(&[0i64, 0, 2], &[3], &wgpu_device); + let s = Tensor::from_slice(&[2.0f32, 4.0, 8.0], &[3], &wgpu_device); + + let mean: Vec = wgpu_client + .scatter_reduce(&d, 0, &i, &s, ScatterReduceOp::Mean, true) + .unwrap() + .to_vec(); + assert_parity_f32(&cpu_mean, &mean, "scatter_reduce_mean_wgpu"); + + let prod: Vec = wgpu_client + .scatter_reduce(&d, 0, &i, &s, ScatterReduceOp::Prod, true) + .unwrap() + .to_vec(); + assert_parity_f32(&cpu_prod, &prod, "scatter_reduce_prod_wgpu"); + }); +} diff --git a/tests/backend_parity/matmul.rs b/tests/backend_parity/matmul.rs index e2551bed..e3355664 100644 --- a/tests/backend_parity/matmul.rs +++ b/tests/backend_parity/matmul.rs @@ -4,10 +4,12 @@ // CPU, CUDA, and WebGPU backends. use numr::ops::MatmulOps; +use numr::runtime::cpu::{CpuClient, CpuDevice, CpuRuntime, ParallelismConfig}; use numr::tensor::Tensor; #[cfg(any(feature = "cuda", feature = "wgpu"))] use crate::backend_parity::helpers::assert_case_parity_f32; +use crate::backend_parity::helpers::assert_parity_f32; #[cfg(feature = "cuda")] use crate::backend_parity::helpers::with_cuda_backend; #[cfg(feature = "wgpu")] @@ -127,3 +129,30 @@ fn test_matmul_vector_parity() { test_matmul_parity(vec![MatmulTest::new(a, vec![1, 4], b, vec![4, 1])]); } + +#[test] +fn test_cpu_matmul_parallelism_config_matches_default() { + let device = CpuDevice::new(); + let default_client = CpuClient::new(device.clone()); + let configured_client = + default_client.with_parallelism(ParallelismConfig::new(Some(1), Some(1024))); + + // Batched case to exercise batch parallelism path. + let a_shape = [6, 24, 16]; + let b_shape = [6, 16, 12]; + let a_numel: usize = a_shape.iter().product(); + let b_numel: usize = b_shape.iter().product(); + let a_data: Vec = (0..a_numel) + .map(|i| (i as f32 * 0.013).sin() + (i as f32 * 0.007).cos()) + .collect(); + let b_data: Vec = (0..b_numel) + .map(|i| (i as f32 * 0.011).cos() - (i as f32 * 0.005).sin()) + .collect(); + + let a = Tensor::::from_slice(&a_data, &a_shape, &device); + let b = Tensor::::from_slice(&b_data, &b_shape, &device); + + let base: Vec = default_client.matmul(&a, &b).unwrap().to_vec(); + let cfg: Vec = configured_client.matmul(&a, &b).unwrap().to_vec(); + assert_parity_f32(&base, &cfg, "cpu_matmul_parallelism_config"); +} diff --git a/tests/backend_parity/matmul_bias.rs b/tests/backend_parity/matmul_bias.rs index 6e9e6a8a..2da812a0 100644 --- a/tests/backend_parity/matmul_bias.rs +++ b/tests/backend_parity/matmul_bias.rs @@ -1,6 +1,7 @@ // Backend parity tests for MatmulOps::matmul_bias use numr::ops::{BinaryOps, MatmulOps}; +use numr::runtime::cpu::{CpuClient, CpuDevice, CpuRuntime, ParallelismConfig}; use numr::tensor::Tensor; use crate::backend_parity::helpers::assert_parity_f32; @@ -107,3 +108,37 @@ fn test_matmul_bias_matches_matmul_plus_bias() { .to_vec(); assert_parity_f32(&fused, &reference, "matmul_bias_matches_reference_cpu"); } + +#[test] +fn test_cpu_matmul_bias_parallelism_config_matches_default() { + let device = CpuDevice::new(); + let default_client = CpuClient::new(device.clone()); + let configured_client = + default_client.with_parallelism(ParallelismConfig::new(Some(1), Some(1024))); + + let a_shape = [4, 20, 16]; + let b_shape = [4, 16, 10]; + let bias_shape = [10]; + let a_numel: usize = a_shape.iter().product(); + let b_numel: usize = b_shape.iter().product(); + let bias_numel: usize = bias_shape.iter().product(); + + let a_data: Vec = (0..a_numel) + .map(|i| (i as f32 * 0.009).sin() + (i as f32 * 0.004).cos()) + .collect(); + let b_data: Vec = (0..b_numel) + .map(|i| (i as f32 * 0.015).cos() - (i as f32 * 0.006).sin()) + .collect(); + let bias_data: Vec = (0..bias_numel).map(|i| (i as f32 * 0.021).sin()).collect(); + + let a = Tensor::::from_slice(&a_data, &a_shape, &device); + let b = Tensor::::from_slice(&b_data, &b_shape, &device); + let bias = Tensor::::from_slice(&bias_data, &bias_shape, &device); + + let base: Vec = default_client.matmul_bias(&a, &b, &bias).unwrap().to_vec(); + let cfg: Vec = configured_client + .matmul_bias(&a, &b, &bias) + .unwrap() + .to_vec(); + assert_parity_f32(&base, &cfg, "cpu_matmul_bias_parallelism_config"); +} diff --git a/tests/backend_parity/reduce.rs b/tests/backend_parity/reduce.rs index 31ee22e1..3898a79d 100644 --- a/tests/backend_parity/reduce.rs +++ b/tests/backend_parity/reduce.rs @@ -5,10 +5,12 @@ use numr::ops::ReduceOps; use numr::runtime::Runtime; +use numr::runtime::cpu::{CpuClient, CpuDevice, CpuRuntime, ParallelismConfig}; use numr::tensor::Tensor; #[cfg(any(feature = "cuda", feature = "wgpu"))] use crate::backend_parity::helpers::assert_case_parity_f32; +use crate::backend_parity::helpers::assert_parity_f32; #[cfg(feature = "cuda")] use crate::backend_parity::helpers::with_cuda_backend; #[cfg(feature = "wgpu")] @@ -128,6 +130,13 @@ fn test_sum_parity() { vec![1], false, ), + // 3D multi-dim reduce + ReduceTest::new( + (1..=24).map(|v| v as f32).collect(), + vec![2, 3, 4], + vec![1, 2], + false, + ), ], ); } @@ -150,6 +159,12 @@ fn test_mean_parity() { vec![1], false, ), + ReduceTest::new( + (1..=24).map(|v| v as f32).collect(), + vec![2, 3, 4], + vec![0, 2], + true, + ), ], ); } @@ -172,6 +187,12 @@ fn test_max_parity() { vec![1], false, ), + ReduceTest::new( + (1..=24).map(|v| v as f32).collect(), + vec![2, 3, 4], + vec![0, 1], + false, + ), ], ); } @@ -194,6 +215,12 @@ fn test_min_parity() { vec![1], false, ), + ReduceTest::new( + (1..=24).map(|v| v as f32).collect(), + vec![2, 3, 4], + vec![0, 1], + false, + ), ], ); } @@ -216,6 +243,12 @@ fn test_prod_parity() { vec![1], false, ), + ReduceTest::new( + vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], + vec![1, 2, 3], + vec![0, 2], + false, + ), ], ); } @@ -243,6 +276,12 @@ fn test_any_parity() { vec![1], false, ), + ReduceTest::new( + vec![0.0, 1.0, 0.0, 0.0, 0.0, 0.0], + vec![1, 2, 3], + vec![0, 2], + false, + ), ], ); } @@ -270,6 +309,53 @@ fn test_all_parity() { vec![1], false, ), + ReduceTest::new( + vec![1.0, 2.0, 3.0, 1.0, 0.0, 3.0], + vec![1, 2, 3], + vec![0, 2], + false, + ), ], ); } + +#[test] +fn test_cpu_reduce_parallelism_config_matches_default() { + let device = CpuDevice::new(); + let default_client = CpuClient::new(device.clone()); + let configured_client = + default_client.with_parallelism(ParallelismConfig::new(Some(1), Some(64))); + + // Large enough to exercise non-last-dim reduction paths where parallel scheduling matters. + let shape = [96, 64, 32]; + let numel: usize = shape.iter().product(); + let data: Vec = (0..numel) + .map(|i| (i as f32 * 0.013).sin() + (i as f32 * 0.007).cos()) + .collect(); + let boolish_data: Vec = (0..numel) + .map(|i| if i % 13 == 0 { 0.0 } else { 1.0 }) + .collect(); + + let a = Tensor::::from_slice(&data, &shape, &device); + let b = Tensor::::from_slice(&boolish_data, &shape, &device); + + let sum_base: Vec = default_client.sum(&a, &[1], false).unwrap().to_vec(); + let sum_cfg: Vec = configured_client.sum(&a, &[1], false).unwrap().to_vec(); + assert_parity_f32(&sum_base, &sum_cfg, "cpu_reduce_parallelism_sum"); + + let mean_base: Vec = default_client.mean(&a, &[1], false).unwrap().to_vec(); + let mean_cfg: Vec = configured_client.mean(&a, &[1], false).unwrap().to_vec(); + assert_parity_f32(&mean_base, &mean_cfg, "cpu_reduce_parallelism_mean"); + + let max_base: Vec = default_client.max(&a, &[1], false).unwrap().to_vec(); + let max_cfg: Vec = configured_client.max(&a, &[1], false).unwrap().to_vec(); + assert_parity_f32(&max_base, &max_cfg, "cpu_reduce_parallelism_max"); + + let prod_base: Vec = default_client.prod(&a, &[1], false).unwrap().to_vec(); + let prod_cfg: Vec = configured_client.prod(&a, &[1], false).unwrap().to_vec(); + assert_parity_f32(&prod_base, &prod_cfg, "cpu_reduce_parallelism_prod"); + + let any_base: Vec = default_client.any(&b, &[1], false).unwrap().to_vec(); + let any_cfg: Vec = configured_client.any(&b, &[1], false).unwrap().to_vec(); + assert_parity_f32(&any_base, &any_cfg, "cpu_reduce_parallelism_any"); +} diff --git a/tests/backend_parity/shape.rs b/tests/backend_parity/shape.rs index 0367309c..495b9618 100644 --- a/tests/backend_parity/shape.rs +++ b/tests/backend_parity/shape.rs @@ -42,6 +42,146 @@ fn test_repeat_on_backends(data: &[f32], shape: &[usize], repeats: &[usize]) { }); } +fn test_cat_on_backends( + a_data: &[f32], + a_shape: &[usize], + b_data: &[f32], + b_shape: &[usize], + dim: isize, +) { + let (cpu_client, cpu_device) = create_cpu_client(); + let a_cpu = Tensor::from_slice(a_data, a_shape, &cpu_device); + let b_cpu = Tensor::from_slice(b_data, b_shape, &cpu_device); + let cpu_result = cpu_client.cat(&[&a_cpu, &b_cpu], dim).unwrap(); + let cpu_data: Vec = cpu_result.to_vec(); + + #[cfg(feature = "cuda")] + with_cuda_backend(|cuda_client, cuda_device| { + let a = Tensor::from_slice(a_data, a_shape, &cuda_device); + let b = Tensor::from_slice(b_data, b_shape, &cuda_device); + let cuda_result = cuda_client.cat(&[&a, &b], dim).unwrap(); + assert_eq!(cpu_result.shape(), cuda_result.shape()); + assert_parity_f32(&cpu_data, &cuda_result.to_vec::(), "cat_cuda"); + }); + + #[cfg(feature = "wgpu")] + with_wgpu_backend(|wgpu_client, wgpu_device| { + let a = Tensor::from_slice(a_data, a_shape, &wgpu_device); + let b = Tensor::from_slice(b_data, b_shape, &wgpu_device); + let wgpu_result = wgpu_client.cat(&[&a, &b], dim).unwrap(); + assert_eq!(cpu_result.shape(), wgpu_result.shape()); + assert_parity_f32(&cpu_data, &wgpu_result.to_vec::(), "cat_wgpu"); + }); +} + +fn test_stack_on_backends( + a_data: &[f32], + a_shape: &[usize], + b_data: &[f32], + b_shape: &[usize], + dim: isize, +) { + let (cpu_client, cpu_device) = create_cpu_client(); + let a_cpu = Tensor::from_slice(a_data, a_shape, &cpu_device); + let b_cpu = Tensor::from_slice(b_data, b_shape, &cpu_device); + let cpu_result = cpu_client.stack(&[&a_cpu, &b_cpu], dim).unwrap(); + let cpu_data: Vec = cpu_result.to_vec(); + + #[cfg(feature = "cuda")] + with_cuda_backend(|cuda_client, cuda_device| { + let a = Tensor::from_slice(a_data, a_shape, &cuda_device); + let b = Tensor::from_slice(b_data, b_shape, &cuda_device); + let cuda_result = cuda_client.stack(&[&a, &b], dim).unwrap(); + assert_eq!(cpu_result.shape(), cuda_result.shape()); + assert_parity_f32(&cpu_data, &cuda_result.to_vec::(), "stack_cuda"); + }); + + #[cfg(feature = "wgpu")] + with_wgpu_backend(|wgpu_client, wgpu_device| { + let a = Tensor::from_slice(a_data, a_shape, &wgpu_device); + let b = Tensor::from_slice(b_data, b_shape, &wgpu_device); + let wgpu_result = wgpu_client.stack(&[&a, &b], dim).unwrap(); + assert_eq!(cpu_result.shape(), wgpu_result.shape()); + assert_parity_f32(&cpu_data, &wgpu_result.to_vec::(), "stack_wgpu"); + }); +} + +fn test_split_on_backends(data: &[f32], shape: &[usize], split_size: usize, dim: isize) { + let (cpu_client, cpu_device) = create_cpu_client(); + let cpu_tensor = Tensor::from_slice(data, shape, &cpu_device); + let cpu_chunks = cpu_client.split(&cpu_tensor, split_size, dim).unwrap(); + let cpu_shapes: Vec> = cpu_chunks.iter().map(|t| t.shape().to_vec()).collect(); + let cpu_data: Vec> = cpu_chunks.iter().map(|t| t.contiguous().to_vec()).collect(); + + #[cfg(feature = "cuda")] + with_cuda_backend(|cuda_client, cuda_device| { + let tensor = Tensor::from_slice(data, shape, &cuda_device); + let chunks = cuda_client.split(&tensor, split_size, dim).unwrap(); + assert_eq!(cpu_chunks.len(), chunks.len()); + for (idx, chunk) in chunks.iter().enumerate() { + assert_eq!(cpu_shapes[idx], chunk.shape().to_vec()); + assert_parity_f32( + &cpu_data[idx], + &chunk.contiguous().to_vec::(), + "split_cuda", + ); + } + }); + + #[cfg(feature = "wgpu")] + with_wgpu_backend(|wgpu_client, wgpu_device| { + let tensor = Tensor::from_slice(data, shape, &wgpu_device); + let chunks = wgpu_client.split(&tensor, split_size, dim).unwrap(); + assert_eq!(cpu_chunks.len(), chunks.len()); + for (idx, chunk) in chunks.iter().enumerate() { + assert_eq!(cpu_shapes[idx], chunk.shape().to_vec()); + assert_parity_f32( + &cpu_data[idx], + &chunk.contiguous().to_vec::(), + "split_wgpu", + ); + } + }); +} + +fn test_chunk_on_backends(data: &[f32], shape: &[usize], chunks: usize, dim: isize) { + let (cpu_client, cpu_device) = create_cpu_client(); + let cpu_tensor = Tensor::from_slice(data, shape, &cpu_device); + let cpu_chunks = cpu_client.chunk(&cpu_tensor, chunks, dim).unwrap(); + let cpu_shapes: Vec> = cpu_chunks.iter().map(|t| t.shape().to_vec()).collect(); + let cpu_data: Vec> = cpu_chunks.iter().map(|t| t.contiguous().to_vec()).collect(); + + #[cfg(feature = "cuda")] + with_cuda_backend(|cuda_client, cuda_device| { + let tensor = Tensor::from_slice(data, shape, &cuda_device); + let got = cuda_client.chunk(&tensor, chunks, dim).unwrap(); + assert_eq!(cpu_chunks.len(), got.len()); + for (idx, chunk) in got.iter().enumerate() { + assert_eq!(cpu_shapes[idx], chunk.shape().to_vec()); + assert_parity_f32( + &cpu_data[idx], + &chunk.contiguous().to_vec::(), + "chunk_cuda", + ); + } + }); + + #[cfg(feature = "wgpu")] + with_wgpu_backend(|wgpu_client, wgpu_device| { + let tensor = Tensor::from_slice(data, shape, &wgpu_device); + let got = wgpu_client.chunk(&tensor, chunks, dim).unwrap(); + assert_eq!(cpu_chunks.len(), got.len()); + for (idx, chunk) in got.iter().enumerate() { + assert_eq!(cpu_shapes[idx], chunk.shape().to_vec()); + assert_parity_f32( + &cpu_data[idx], + &chunk.contiguous().to_vec::(), + "chunk_wgpu", + ); + } + }); +} + fn test_pad_on_backends(data: &[f32], shape: &[usize], padding: &[usize], value: f64) { let (cpu_client, cpu_device) = create_cpu_client(); let cpu_tensor = Tensor::from_slice(data, shape, &cpu_device); @@ -88,6 +228,79 @@ fn test_roll_on_backends(data: &[f32], shape: &[usize], shift: isize, dim: isize }); } +fn test_unfold_on_backends(data: &[f32], shape: &[usize], dim: isize, size: usize, step: usize) { + let (cpu_client, cpu_device) = create_cpu_client(); + let cpu_tensor = Tensor::from_slice(data, shape, &cpu_device); + let cpu_result = cpu_client.unfold(&cpu_tensor, dim, size, step).unwrap(); + let cpu_data: Vec = cpu_result.contiguous().to_vec(); + + #[cfg(feature = "cuda")] + with_cuda_backend(|cuda_client, cuda_device| { + let cuda_tensor = Tensor::from_slice(data, shape, &cuda_device); + let cuda_result = cuda_client.unfold(&cuda_tensor, dim, size, step).unwrap(); + assert_eq!(cpu_result.shape(), cuda_result.shape()); + assert_parity_f32( + &cpu_data, + &cuda_result.contiguous().to_vec::(), + "unfold_cuda", + ); + }); + + #[cfg(feature = "wgpu")] + with_wgpu_backend(|wgpu_client, wgpu_device| { + let wgpu_tensor = Tensor::from_slice(data, shape, &wgpu_device); + let wgpu_result = wgpu_client.unfold(&wgpu_tensor, dim, size, step).unwrap(); + assert_eq!(cpu_result.shape(), wgpu_result.shape()); + assert_parity_f32( + &cpu_data, + &wgpu_result.contiguous().to_vec::(), + "unfold_wgpu", + ); + }); +} + +fn test_repeat_interleave_on_backends( + data: &[f32], + shape: &[usize], + repeats: usize, + dim: Option, +) { + let (cpu_client, cpu_device) = create_cpu_client(); + let cpu_tensor = Tensor::from_slice(data, shape, &cpu_device); + let cpu_result = cpu_client + .repeat_interleave(&cpu_tensor, repeats, dim) + .unwrap(); + let cpu_data: Vec = cpu_result.to_vec(); + + #[cfg(feature = "cuda")] + with_cuda_backend(|cuda_client, cuda_device| { + let cuda_tensor = Tensor::from_slice(data, shape, &cuda_device); + let cuda_result = cuda_client + .repeat_interleave(&cuda_tensor, repeats, dim) + .unwrap(); + assert_eq!(cpu_result.shape(), cuda_result.shape()); + assert_parity_f32( + &cpu_data, + &cuda_result.to_vec::(), + "repeat_interleave_cuda", + ); + }); + + #[cfg(feature = "wgpu")] + with_wgpu_backend(|wgpu_client, wgpu_device| { + let wgpu_tensor = Tensor::from_slice(data, shape, &wgpu_device); + let wgpu_result = wgpu_client + .repeat_interleave(&wgpu_tensor, repeats, dim) + .unwrap(); + assert_eq!(cpu_result.shape(), wgpu_result.shape()); + assert_parity_f32( + &cpu_data, + &wgpu_result.to_vec::(), + "repeat_interleave_wgpu", + ); + }); +} + fn test_flip_on_backends(data: &[f32], shape: &[usize], dim: isize) { use numr::runtime::cpu::{CpuDevice, CpuRuntime}; let cpu_device = CpuDevice::new(); @@ -126,6 +339,32 @@ fn test_flip_on_backends(data: &[f32], shape: &[usize], dim: isize) { // Shape Operation Parity Tests // ============================================================================ +#[test] +fn test_cat_parity_negative_dim() { + let a = [1.0f32, 2.0, 3.0, 4.0]; + let b = [10.0f32, 20.0]; + test_cat_on_backends(&a, &[2, 2], &b, &[2, 1], -1); +} + +#[test] +fn test_stack_parity_negative_dim() { + let a = [1.0f32, 2.0, 3.0, 4.0]; + let b = [10.0f32, 20.0, 30.0, 40.0]; + test_stack_on_backends(&a, &[2, 2], &b, &[2, 2], -1); +} + +#[test] +fn test_split_parity_negative_dim() { + let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]; + test_split_on_backends(&data, &[2, 5], 2, -1); +} + +#[test] +fn test_chunk_parity_negative_dim() { + let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]; + test_chunk_on_backends(&data, &[2, 5], 3, -1); +} + #[test] fn test_repeat_parity() { let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; @@ -145,8 +384,56 @@ fn test_roll_parity() { test_roll_on_backends(&data, &[2, 3], 2, 1); } +#[test] +fn test_roll_parity_negative_dim() { + let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; + test_roll_on_backends(&data, &[2, 3], -1, -1); +} + #[test] fn test_flip_parity() { let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; test_flip_on_backends(&data, &[2, 3], 1); } + +#[test] +fn test_flip_parity_negative_dim() { + let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; + test_flip_on_backends(&data, &[2, 3], -1); +} + +#[test] +fn test_unfold_parity() { + let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; + test_unfold_on_backends(&data, &[2, 3], 1, 2, 1); +} + +#[test] +fn test_unfold_parity_dim0() { + let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; + test_unfold_on_backends(&data, &[2, 3], 0, 2, 1); +} + +#[test] +fn test_unfold_parity_negative_dim() { + let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; + test_unfold_on_backends(&data, &[2, 3], -1, 2, 1); +} + +#[test] +fn test_repeat_interleave_parity() { + let data = [1.0f32, 2.0, 3.0, 4.0]; + test_repeat_interleave_on_backends(&data, &[2, 2], 2, Some(1)); +} + +#[test] +fn test_repeat_interleave_parity_negative_dim() { + let data = [1.0f32, 2.0, 3.0, 4.0]; + test_repeat_interleave_on_backends(&data, &[2, 2], 2, Some(-1)); +} + +#[test] +fn test_repeat_interleave_parity_flattened() { + let data = [1.0f32, 2.0, 3.0, 4.0]; + test_repeat_interleave_on_backends(&data, &[2, 2], 2, None); +}