Skip to main content

etensor_core/backends/cpu/
binary.rs

1//! Binary element-wise operations for the CPU execution backend.
2//!
3//! These operations are memory-bandwidth-bound (just read + write, minimal math).
4//! Parallelism via Rayon is NOT used because the memory bus is the bottleneck,
5//! not the CPU. Adding thread wake-up overhead (~20µs) hurts performance
6//! at all practical tensor sizes. The single-threaded vectorized loop already
7//! saturates the memory bus on modern CPUs.
8
9use crate::tensor::Tensor;
10use crate::buffer::Buffer;
11use crate::errors::EtensorResult;
12use crate::backends::traits::Backend;
13
14/// The execution driver for standard host CPU memory.
15pub struct CpuBackend;
16
17impl Backend for CpuBackend {
18    fn add(a: &Tensor, b: &Tensor) -> EtensorResult<Tensor> {
19        let slice_a = a.data.as_f32_slice()?;
20        let slice_b = b.data.as_f32_slice()?;
21
22        // Single-threaded vectorized loop. LLVM auto-vectorizes this to SIMD
23        // (SSE2/AVX2) which already saturates memory bandwidth on one core.
24        let out_vec: Vec<f32> = slice_a.iter().zip(slice_b).map(|(x, y)| x + y).collect();
25
26        Ok(Tensor::new(
27            Buffer::from_f32_vec(out_vec),
28            a.shape.clone(),
29            a.device,
30            a.dtype,
31            false, // Gradients are exclusively managed by the Dispatcher, not the hardware kernel.
32        ))
33    }
34
35    fn mul(a: &Tensor, b: &Tensor) -> EtensorResult<Tensor> {
36        let slice_a = a.data.as_f32_slice()?;
37        let slice_b = b.data.as_f32_slice()?;
38
39        let out_vec: Vec<f32> = slice_a.iter().zip(slice_b).map(|(x, y)| x * y).collect();
40
41        Ok(Tensor::new(
42            Buffer::from_f32_vec(out_vec),
43            a.shape.clone(),
44            a.device,
45            a.dtype,
46            false,
47        ))
48    }
49
50    fn matmul(a: &Tensor, b: &Tensor) -> EtensorResult<Tensor> {
51        crate::backends::cpu::matmul::matmul_forward(a, b)
52    }
53
54    fn sum_all(a: &Tensor) -> EtensorResult<Tensor> {
55        crate::backends::cpu::reduce::sum_all(a)
56    }
57
58    fn mean_all(a: &Tensor) -> EtensorResult<Tensor> {
59        crate::backends::cpu::reduce::mean_all(a)
60    }
61
62    fn max_all(a: &Tensor) -> EtensorResult<Tensor> {
63        crate::backends::cpu::reduce::max_all(a)
64    }
65
66    fn relu(a: &Tensor) -> EtensorResult<Tensor> {
67        crate::backends::cpu::unary::relu_forward(a)
68    }
69
70    fn sigmoid(a: &Tensor) -> EtensorResult<Tensor> {
71        crate::backends::cpu::unary::sigmoid_forward(a)
72    }
73
74    fn add_relu(a: &Tensor, b: &Tensor) -> EtensorResult<Tensor> {
75        crate::backends::cpu::fusion::add_relu_forward(a, b)
76    }
77
78    fn linear(x: &Tensor, w: &Tensor, b: &Tensor) -> EtensorResult<Tensor> {
79        crate::backends::cpu::fusion::linear_forward(x, w, b)
80    }
81}
82
83// =====================================================================
84// UNIT TESTS
85// =====================================================================
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use crate::errors::EtensorError;
90    use crate::device::Device;
91    use crate::dtypes::DType;
92    use crate::shape::Shape;
93
94    // Helper to generate purely physical test tensors
95    fn make_test_tensor(data: Vec<f32>) -> Tensor {
96        let len = data.len();
97        Tensor::new(
98            Buffer::from_f32_vec(data),
99            Shape::new(vec![len]),
100            Device::Cpu,
101            DType::F32,
102            false,
103        )
104    }
105
106    #[test]
107    fn test_cpu_backend_add() {
108        let a = make_test_tensor(vec![1.5, 2.5, 3.5]);
109        let b = make_test_tensor(vec![2.0, 3.0, 4.0]);
110        
111        let c = CpuBackend::add(&a, &b).unwrap();
112        let slice = c.data.as_f32_slice().unwrap();
113        
114        assert_eq!(slice, &[3.5, 5.5, 7.5]);
115    }
116
117    #[test]
118    fn test_cpu_backend_mul() {
119        let a = make_test_tensor(vec![2.0, 4.0, 6.0]);
120        let b = make_test_tensor(vec![3.0, 0.5, 10.0]);
121        
122        let c = CpuBackend::mul(&a, &b).unwrap();
123        let slice = c.data.as_f32_slice().unwrap();
124        
125        assert_eq!(slice, &[6.0, 2.0, 60.0]);
126    }
127
128    #[test]
129    fn test_cpu_backend_rejects_gpu_tensors() {
130        let shape = Shape::new(vec![1]);
131        
132        let a = Tensor::new(
133            Buffer::new_cpu_zeros(1, DType::I32), 
134            shape.clone(), 
135            Device::Cpu, 
136            DType::I32, 
137            false
138        );
139        let b = Tensor::new(
140            Buffer::new_cpu_zeros(1, DType::I32), 
141            shape, 
142            Device::Cpu, 
143            DType::I32, 
144            false
145        );
146
147        let result = CpuBackend::add(&a, &b);
148        assert!(result.is_err());
149        
150        if let Err(EtensorError::DTypeMismatch { .. }) = result {
151            // Success: The backend protected itself from invalid memory layouts
152        } else {
153            panic!("Expected DTypeMismatch error from the CPU Kernel!");
154        }
155    }
156}