etensor_core/backends/cpu/
reduce.rs

1//! High-performance CPU memory reduction kernels.
2//! 
3//! Reductions are bandwidth-bound (read N elements, write 1 scalar).
4//! Single-threaded loops with LLVM auto-vectorization already saturate
5//! the memory bus. Rayon is not used.
6
7use crate::tensor::Tensor;
8use crate::buffer::Buffer;
9use crate::shape::Shape;
10use crate::device::Device;
11use crate::errors::{EtensorError, EtensorResult};
12
13/// Executes a global sum reduction, collapsing the entire tensor into a single scalar value.
14pub fn sum_all(a: &Tensor) -> EtensorResult<Tensor> {
15    let slice = a.data.as_f32_slice()?;
16    
17    let total_sum: f32 = slice.iter().sum();
18
19    let out_shape = Shape::new(vec![1]);
20    
21    Ok(Tensor::new(
22        Buffer::from_f32_vec(vec![total_sum]),
23        out_shape,
24        Device::Cpu,
25        a.dtype,
26        false, // Gradients are exclusively managed by the Dispatcher.
27    ))
28}
29
30/// Executes a global mean reduction, calculating the average of all elements.
31pub fn mean_all(a: &Tensor) -> EtensorResult<Tensor> {
32    let slice = a.data.as_f32_slice()?;
33    
34    if slice.is_empty() {
35        return Err(EtensorError::InternalError(
36            "Cannot calculate the mean of an empty tensor.".to_string(),
37        ));
38    }
39
40    let total_sum: f32 = slice.iter().sum();
41    let mean = total_sum / (slice.len() as f32);
42
43    let out_shape = Shape::new(vec![1]);
44    
45    Ok(Tensor::new(
46        Buffer::from_f32_vec(vec![mean]),
47        out_shape,
48        Device::Cpu,
49        a.dtype,
50        false,
51    ))
52}
53
54/// Executes a global max reduction, finding the largest single value in the tensor.
55pub fn max_all(a: &Tensor) -> EtensorResult<Tensor> {
56    let slice = a.data.as_f32_slice()?;
57    
58    if slice.is_empty() {
59        return Err(EtensorError::InternalError(
60            "Cannot calculate the max of an empty tensor.".to_string(),
61        ));
62    }
63
64    let max_val = slice.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
65
66    let out_shape = Shape::new(vec![1]);
67    
68    Ok(Tensor::new(
69        Buffer::from_f32_vec(vec![max_val]),
70        out_shape,
71        Device::Cpu,
72        a.dtype,
73        false,
74    ))
75}
76
77// =====================================================================
78// UNIT TESTS
79// =====================================================================
80#[cfg(test)]
81mod tests {
82    use super::*;
83    use crate::dtypes::DType;
84
85    fn make_test_tensor(data: Vec<f32>, dims: Vec<usize>) -> Tensor {
86        Tensor::new(
87            Buffer::from_f32_vec(data),
88            Shape::new(dims),
89            Device::Cpu,
90            DType::F32,
91            false,
92        )
93    }
94
95    #[test]
96    fn test_cpu_reduce_sum_all() {
97        let a = make_test_tensor(vec![1.0, 2.0, 3.0, 4.0], vec![2, 2]);
98        
99        let c = sum_all(&a).unwrap();
100        let slice = c.data.as_f32_slice().unwrap();
101
102        assert_eq!(c.shape.dims, vec![1]);
103        assert_eq!(slice, &[10.0]);
104    }
105
106    #[test]
107    fn test_cpu_reduce_mean_all() {
108        let a = make_test_tensor(vec![2.0, 4.0, 6.0, 8.0], vec![4]);
109        
110        let c = mean_all(&a).unwrap();
111        let slice = c.data.as_f32_slice().unwrap();
112
113        assert_eq!(c.shape.dims, vec![1]);
114        assert_eq!(slice, &[5.0]);
115    }
116
117    #[test]
118    fn test_cpu_reduce_max_all() {
119        let a = make_test_tensor(vec![-5.0, 12.0, 3.0, 42.0, 0.0, -100.0], vec![2, 3, 1]);
120        
121        let c = max_all(&a).unwrap();
122        let slice = c.data.as_f32_slice().unwrap();
123
124        assert_eq!(c.shape.dims, vec![1]);
125        assert_eq!(slice, &[42.0]);
126    }
127}
etensor_core/backends/cpu/reduce.rs

etensor_core/backends/cpu/
reduce.rs
