etensor_core/backends/cpu/
unary.rs

1//! High-performance CPU Unary operation kernels.
2//! 
3//! ReLU is bandwidth-bound (simple comparison) → single-threaded, LLVM auto-vectorizes.
4//! Sigmoid is compute-bound (exp() is expensive) → uses Rayon for large tensors.
5
6use crate::tensor::Tensor;
7use crate::buffer::Buffer;
8use crate::device::Device;
9use crate::errors::EtensorResult;
10use rayon::prelude::*;
11
12/// Sigmoid is compute-bound (exp() ~20 cycles vs add's ~1 cycle), so parallelism helps.
13/// The threshold is set high enough that Rayon's ~20µs overhead is amortized.
14const RAYON_THRESHOLD_COMPUTE: usize = 131_072; // 128K elements
15
16/// Executes the Rectified Linear Unit (ReLU) activation function: f(x) = max(0, x)
17/// 
18/// Bandwidth-bound: a single max() comparison per element. Single-threaded 
19/// SIMD already saturates the memory bus.
20pub fn relu_forward(a: &Tensor) -> EtensorResult<Tensor> {
21    let slice = a.data.as_f32_slice()?;
22    
23    let out_vec: Vec<f32> = slice.iter().map(|&x| x.max(0.0)).collect();
24
25    Ok(Tensor::new(
26        Buffer::from_f32_vec(out_vec),
27        a.shape.clone(),
28        Device::Cpu,
29        a.dtype,
30        false, // Gradients are exclusively managed by the Dispatcher.
31    ))
32}
33
34/// Executes the Sigmoid activation function: f(x) = 1 / (1 + exp(-x))
35/// 
36/// Compute-bound: exp() is ~20x more expensive than add/mul per element.
37/// Rayon parallelism provides genuine speedup at large sizes.
38pub fn sigmoid_forward(a: &Tensor) -> EtensorResult<Tensor> {
39    let slice = a.data.as_f32_slice()?;
40    
41    let out_vec: Vec<f32> = if slice.len() < RAYON_THRESHOLD_COMPUTE {
42        slice.iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect()
43    } else {
44        slice.par_iter().map(|&x| 1.0 / (1.0 + (-x).exp())).collect()
45    };
46
47    Ok(Tensor::new(
48        Buffer::from_f32_vec(out_vec),
49        a.shape.clone(),
50        Device::Cpu,
51        a.dtype,
52        false,
53    ))
54}
55
56// =====================================================================
57// UNIT TESTS
58// =====================================================================
59#[cfg(test)]
60mod tests {
61    use super::*;
62    use crate::shape::Shape;
63    use crate::dtypes::DType;
64
65    // Helper to generate a test matrix
66    fn make_test_tensor(data: Vec<f32>) -> Tensor {
67        let len = data.len();
68        Tensor::new(
69            Buffer::from_f32_vec(data),
70            Shape::new(vec![len]),
71            Device::Cpu,
72            DType::F32,
73            false,
74        )
75    }
76
77    #[test]
78    fn test_cpu_relu() {
79        let a = make_test_tensor(vec![-5.0, 0.0, 3.14, -0.01, 42.0]);
80        
81        let c = relu_forward(&a).unwrap();
82        let slice = c.data.as_f32_slice().unwrap();
83
84        assert_eq!(slice, &[0.0, 0.0, 3.14, 0.0, 42.0]);
85    }
86
87    #[test]
88    fn test_cpu_sigmoid() {
89        let a = make_test_tensor(vec![0.0, 100.0, -100.0]);
90        
91        let c = sigmoid_forward(&a).unwrap();
92        let slice = c.data.as_f32_slice().unwrap();
93
94        assert_eq!(slice[0], 0.5);
95        assert!(slice[1] > 0.999);
96        assert!(slice[2] < 0.001);
97    }
98}
etensor_core/backends/cpu/unary.rs

etensor_core/backends/cpu/
unary.rs
