Skip to main content

tenflowers_core/ops/binary/
operations.rs

1//! Concrete Binary Operation Implementations
2//!
3//! This module provides concrete implementations of all binary operations with
4//! ultra-performance optimizations including SIMD acceleration and parallel processing.
5
6use super::core::{get_binary_op_registry, BinaryOp, OpComplexity};
7use super::simd::simd_f32_ops;
8use crate::{Result, TensorError};
9use rayon::prelude::*;
10use scirs2_core::numeric::Zero;
11use std::ops::{Add as StdAdd, Div as StdDiv, Mul as StdMul, Sub as StdSub};
12
13/// Ultra-performance addition operation with SIMD and parallel support
14#[derive(Clone)]
15pub struct AddOp;
16
17impl<T: StdAdd<Output = T> + Clone + Send + Sync + 'static> BinaryOp<T> for AddOp {
18    #[inline]
19    fn apply(&self, a: T, b: T) -> T {
20        a + b
21    }
22
23    #[inline]
24    fn name(&self) -> &str {
25        "Add"
26    }
27
28    fn apply_slice(&self, a: &[T], b: &[T], output: &mut [T]) -> Result<()> {
29        if a.len() != b.len() || a.len() != output.len() {
30            return Err(TensorError::invalid_argument(
31                "Slice length mismatch for Add operation".to_string(),
32            ));
33        }
34
35        // Special case for f32 - use SIMD optimization
36        if std::any::type_name::<T>() == "f32" {
37            let a_f32 = unsafe { std::slice::from_raw_parts(a.as_ptr() as *const f32, a.len()) };
38            let b_f32 = unsafe { std::slice::from_raw_parts(b.as_ptr() as *const f32, b.len()) };
39            let output_f32 = unsafe {
40                std::slice::from_raw_parts_mut(output.as_mut_ptr() as *mut f32, output.len())
41            };
42            return simd_f32_ops::simd_add_f32(a_f32, b_f32, output_f32).map_err(|_| {
43                TensorError::invalid_argument("SIMD Add operation failed".to_string())
44            });
45        }
46
47        // Use parallel processing for large arrays
48        if a.len() >= 10000 {
49            get_binary_op_registry().record_parallel_usage();
50            output
51                .par_iter_mut()
52                .zip(a.par_iter().zip(b.par_iter()))
53                .for_each(|(out, (a_val, b_val))| {
54                    *out = a_val.clone() + b_val.clone();
55                });
56        } else {
57            // Sequential for small arrays
58            for i in 0..a.len() {
59                output[i] = a[i].clone() + b[i].clone();
60            }
61        }
62        Ok(())
63    }
64
65    fn supports_simd(&self) -> bool {
66        true
67    }
68    fn supports_gpu(&self) -> bool {
69        true
70    }
71    fn complexity(&self) -> OpComplexity {
72        OpComplexity::Simple
73    }
74    fn is_associative(&self) -> bool {
75        true
76    }
77    fn is_commutative(&self) -> bool {
78        true
79    }
80}
81
82/// Ultra-performance subtraction operation with SIMD and parallel support
83#[derive(Clone)]
84pub struct SubOp;
85
86impl<T: StdSub<Output = T> + Clone + Send + Sync + 'static> BinaryOp<T> for SubOp {
87    #[inline]
88    fn apply(&self, a: T, b: T) -> T {
89        a - b
90    }
91
92    #[inline]
93    fn name(&self) -> &str {
94        "Sub"
95    }
96
97    fn apply_slice(&self, a: &[T], b: &[T], output: &mut [T]) -> Result<()> {
98        if a.len() != b.len() || a.len() != output.len() {
99            return Err(TensorError::invalid_argument(
100                "Slice length mismatch for Sub operation".to_string(),
101            ));
102        }
103
104        // Use parallel processing for large arrays
105        if a.len() >= 10000 {
106            get_binary_op_registry().record_parallel_usage();
107            output
108                .par_iter_mut()
109                .zip(a.par_iter().zip(b.par_iter()))
110                .for_each(|(out, (a_val, b_val))| {
111                    *out = a_val.clone() - b_val.clone();
112                });
113        } else {
114            // Sequential for small arrays
115            for i in 0..a.len() {
116                output[i] = a[i].clone() - b[i].clone();
117            }
118        }
119        Ok(())
120    }
121
122    fn supports_simd(&self) -> bool {
123        true
124    }
125    fn supports_gpu(&self) -> bool {
126        true
127    }
128    fn complexity(&self) -> OpComplexity {
129        OpComplexity::Simple
130    }
131    fn is_associative(&self) -> bool {
132        false
133    } // Subtraction is not associative
134    fn is_commutative(&self) -> bool {
135        false
136    } // Subtraction is not commutative
137}
138
139/// Ultra-performance multiplication operation with SIMD and parallel support
140#[derive(Clone)]
141pub struct MulOp;
142
143impl<T: StdMul<Output = T> + Clone + Send + Sync + 'static> BinaryOp<T> for MulOp {
144    #[inline]
145    fn apply(&self, a: T, b: T) -> T {
146        a * b
147    }
148    #[inline]
149    fn name(&self) -> &str {
150        "Mul"
151    }
152
153    fn apply_slice(&self, a: &[T], b: &[T], output: &mut [T]) -> Result<()> {
154        if a.len() != b.len() || a.len() != output.len() {
155            return Err(TensorError::invalid_argument(
156                "Slice length mismatch for Mul operation".to_string(),
157            ));
158        }
159
160        // Special case for f32 - use SIMD optimization
161        if std::any::type_name::<T>() == "f32" {
162            let a_f32 = unsafe { std::slice::from_raw_parts(a.as_ptr() as *const f32, a.len()) };
163            let b_f32 = unsafe { std::slice::from_raw_parts(b.as_ptr() as *const f32, b.len()) };
164            let output_f32 = unsafe {
165                std::slice::from_raw_parts_mut(output.as_mut_ptr() as *mut f32, output.len())
166            };
167            return simd_f32_ops::simd_mul_f32(a_f32, b_f32, output_f32).map_err(|_| {
168                TensorError::invalid_argument("SIMD Mul operation failed".to_string())
169            });
170        }
171
172        // Use parallel processing for large arrays
173        if a.len() >= 8000 {
174            get_binary_op_registry().record_parallel_usage();
175            output
176                .par_iter_mut()
177                .zip(a.par_iter().zip(b.par_iter()))
178                .for_each(|(out, (a_val, b_val))| {
179                    *out = a_val.clone() * b_val.clone();
180                });
181        } else {
182            // Sequential for small arrays
183            for i in 0..a.len() {
184                output[i] = a[i].clone() * b[i].clone();
185            }
186        }
187        Ok(())
188    }
189
190    fn supports_simd(&self) -> bool {
191        true
192    }
193    fn supports_gpu(&self) -> bool {
194        true
195    }
196    fn complexity(&self) -> OpComplexity {
197        OpComplexity::Simple
198    }
199    fn is_associative(&self) -> bool {
200        true
201    }
202    fn is_commutative(&self) -> bool {
203        true
204    }
205}
206
207/// Ultra-performance division operation with SIMD and parallel support
208#[derive(Clone)]
209pub struct DivOp;
210
211impl<T: StdDiv<Output = T> + Clone + Send + Sync + 'static> BinaryOp<T> for DivOp {
212    #[inline]
213    fn apply(&self, a: T, b: T) -> T {
214        a / b
215    }
216
217    #[inline]
218    fn name(&self) -> &str {
219        "Div"
220    }
221
222    fn apply_slice(&self, a: &[T], b: &[T], output: &mut [T]) -> Result<()> {
223        if a.len() != b.len() || a.len() != output.len() {
224            return Err(TensorError::invalid_argument(
225                "Slice length mismatch for Div operation".to_string(),
226            ));
227        }
228
229        // Use parallel processing for large arrays (division is more expensive)
230        if a.len() >= 5000 {
231            get_binary_op_registry().record_parallel_usage();
232            output
233                .par_iter_mut()
234                .zip(a.par_iter().zip(b.par_iter()))
235                .for_each(|(out, (a_val, b_val))| {
236                    *out = a_val.clone() / b_val.clone();
237                });
238        } else {
239            // Sequential for small arrays
240            for i in 0..a.len() {
241                output[i] = a[i].clone() / b[i].clone();
242            }
243        }
244        Ok(())
245    }
246
247    fn supports_simd(&self) -> bool {
248        true
249    }
250    fn supports_gpu(&self) -> bool {
251        true
252    }
253    fn complexity(&self) -> OpComplexity {
254        OpComplexity::Moderate
255    } // Division is more expensive
256    fn is_associative(&self) -> bool {
257        false
258    } // Division is not associative
259    fn is_commutative(&self) -> bool {
260        false
261    } // Division is not commutative
262}
263
264/// Power operation
265#[derive(Clone)]
266pub struct PowOp;
267impl<T: scirs2_core::num_traits::Float> BinaryOp<T> for PowOp {
268    #[inline]
269    fn apply(&self, a: T, b: T) -> T {
270        a.powf(b)
271    }
272    #[inline]
273    fn name(&self) -> &str {
274        "Pow"
275    }
276
277    fn complexity(&self) -> OpComplexity {
278        OpComplexity::Complex
279    }
280    fn supports_gpu(&self) -> bool {
281        true
282    }
283}
284
285/// Element-wise minimum operation
286#[derive(Clone)]
287pub struct MinOp;
288impl<T: PartialOrd + Clone> BinaryOp<T> for MinOp {
289    #[inline]
290    fn apply(&self, a: T, b: T) -> T {
291        if a <= b {
292            a
293        } else {
294            b
295        }
296    }
297    #[inline]
298    fn name(&self) -> &str {
299        "Min"
300    }
301
302    fn complexity(&self) -> OpComplexity {
303        OpComplexity::Moderate
304    }
305    fn supports_simd(&self) -> bool {
306        true
307    }
308    fn supports_gpu(&self) -> bool {
309        true
310    }
311    fn is_associative(&self) -> bool {
312        true
313    }
314    fn is_commutative(&self) -> bool {
315        true
316    }
317}
318
319/// Element-wise maximum operation
320#[derive(Clone)]
321pub struct MaxOp;
322impl<T: PartialOrd + Clone> BinaryOp<T> for MaxOp {
323    #[inline]
324    fn apply(&self, a: T, b: T) -> T {
325        if a >= b {
326            a
327        } else {
328            b
329        }
330    }
331    #[inline]
332    fn name(&self) -> &str {
333        "Max"
334    }
335
336    fn complexity(&self) -> OpComplexity {
337        OpComplexity::Moderate
338    }
339    fn supports_simd(&self) -> bool {
340        true
341    }
342    fn supports_gpu(&self) -> bool {
343        true
344    }
345    fn is_associative(&self) -> bool {
346        true
347    }
348    fn is_commutative(&self) -> bool {
349        true
350    }
351}
352
353/// PReLU operation: PReLU(x, alpha) = x if x > 0, else alpha * x
354#[derive(Clone)]
355pub struct PReLUOp;
356impl<T> BinaryOp<T> for PReLUOp
357where
358    T: scirs2_core::num_traits::Float + PartialOrd + Zero + StdMul<Output = T>,
359{
360    #[inline]
361    fn apply(&self, x: T, alpha: T) -> T {
362        if x > T::zero() {
363            x
364        } else {
365            alpha * x
366        }
367    }
368    #[inline]
369    fn name(&self) -> &str {
370        "PReLU"
371    }
372
373    fn complexity(&self) -> OpComplexity {
374        OpComplexity::Moderate
375    }
376    fn supports_gpu(&self) -> bool {
377        true
378    }
379}