tenflowers_core/ops/binary/
operations.rs1use 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#[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 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 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 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#[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 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 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 } fn is_commutative(&self) -> bool {
135 false
136 } }
138
139#[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 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 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 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#[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 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 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 } fn is_associative(&self) -> bool {
257 false
258 } fn is_commutative(&self) -> bool {
260 false
261 } }
263
264#[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#[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#[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#[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}