tenflowers_core/simd/
basic_ops.rs

1//! Basic SIMD Operations
2//!
3//! This module provides optimized element-wise operations using SIMD instructions
4//! for fundamental tensor operations like addition, multiplication, and fused multiply-add.
5
6use crate::error::ErrorContext;
7use crate::{Result, TensorError};
8
9/// Basic SIMD-optimized element-wise operations
10pub struct BasicOps;
11
12impl BasicOps {
13    /// Fast inline element-wise addition without bounds checking (for hot paths)
14    #[inline(always)]
15    pub fn add_f32_unchecked(a: &[f32], b: &[f32], result: &mut [f32]) {
16        let len = a.len();
17
18        // For very small arrays, use simple vectorizable loop
19        if len < 32 {
20            for i in 0..len {
21                unsafe {
22                    *result.get_unchecked_mut(i) = a.get_unchecked(i) + b.get_unchecked(i);
23                }
24            }
25            return;
26        }
27
28        // For larger arrays, use optimal chunk size for SIMD (8 elements = 256 bits for AVX)
29        let chunks = len / 8;
30        let remainder = len % 8;
31
32        // Process 8 elements at a time - optimal for most SIMD instruction sets
33        for chunk in 0..chunks {
34            let base = chunk * 8;
35
36            // Compact unrolling that's more likely to vectorize efficiently
37            unsafe {
38                *result.get_unchecked_mut(base) = a.get_unchecked(base) + b.get_unchecked(base);
39                *result.get_unchecked_mut(base + 1) =
40                    a.get_unchecked(base + 1) + b.get_unchecked(base + 1);
41                *result.get_unchecked_mut(base + 2) =
42                    a.get_unchecked(base + 2) + b.get_unchecked(base + 2);
43                *result.get_unchecked_mut(base + 3) =
44                    a.get_unchecked(base + 3) + b.get_unchecked(base + 3);
45                *result.get_unchecked_mut(base + 4) =
46                    a.get_unchecked(base + 4) + b.get_unchecked(base + 4);
47                *result.get_unchecked_mut(base + 5) =
48                    a.get_unchecked(base + 5) + b.get_unchecked(base + 5);
49                *result.get_unchecked_mut(base + 6) =
50                    a.get_unchecked(base + 6) + b.get_unchecked(base + 6);
51                *result.get_unchecked_mut(base + 7) =
52                    a.get_unchecked(base + 7) + b.get_unchecked(base + 7);
53            }
54        }
55
56        // Handle remaining elements
57        let remainder_start = chunks * 8;
58        for i in 0..remainder {
59            unsafe {
60                *result.get_unchecked_mut(remainder_start + i) =
61                    a.get_unchecked(remainder_start + i) + b.get_unchecked(remainder_start + i);
62            }
63        }
64    }
65
66    /// Element-wise addition with SciRS2-Core SIMD auto-vectorization (safe version)
67    pub fn add_f32_optimized(a: &[f32], b: &[f32], result: &mut [f32]) -> Result<()> {
68        if a.len() != b.len() || a.len() != result.len() {
69            return Err(TensorError::ShapeMismatch {
70                operation: "SIMD add_f32".to_string(),
71                expected: format!("arrays of length {}", a.len()),
72                got: format!("a: {}, b: {}, result: {}", a.len(), b.len(), result.len()),
73                context: Some(Box::new(ErrorContext::new())),
74            });
75        }
76
77        // Use basic SIMD when available (advanced SIMD features pending in scirs2_core)
78        #[cfg(feature = "simd")]
79        {
80            // For now, use basic vectorization - advanced SIMD will be added when available
81            if a.len() >= 64 {
82                // Fallback to optimized manual implementation
83            }
84        }
85
86        Self::add_f32_unchecked(a, b, result);
87        Ok(())
88    }
89
90    /// Fast inline element-wise multiplication without bounds checking (for hot paths)
91    #[inline(always)]
92    pub fn mul_f32_unchecked(a: &[f32], b: &[f32], result: &mut [f32]) {
93        let len = a.len();
94
95        // For very small arrays, use simple vectorizable loop
96        if len < 32 {
97            for i in 0..len {
98                unsafe {
99                    *result.get_unchecked_mut(i) = a.get_unchecked(i) * b.get_unchecked(i);
100                }
101            }
102            return;
103        }
104
105        // For larger arrays, use optimal chunk size for SIMD (8 elements = 256 bits for AVX)
106        let chunks = len / 8;
107        let remainder = len % 8;
108
109        // Process 8 elements at a time - optimal for most SIMD instruction sets
110        for chunk in 0..chunks {
111            let base = chunk * 8;
112
113            unsafe {
114                *result.get_unchecked_mut(base) = a.get_unchecked(base) * b.get_unchecked(base);
115                *result.get_unchecked_mut(base + 1) =
116                    a.get_unchecked(base + 1) * b.get_unchecked(base + 1);
117                *result.get_unchecked_mut(base + 2) =
118                    a.get_unchecked(base + 2) * b.get_unchecked(base + 2);
119                *result.get_unchecked_mut(base + 3) =
120                    a.get_unchecked(base + 3) * b.get_unchecked(base + 3);
121                *result.get_unchecked_mut(base + 4) =
122                    a.get_unchecked(base + 4) * b.get_unchecked(base + 4);
123                *result.get_unchecked_mut(base + 5) =
124                    a.get_unchecked(base + 5) * b.get_unchecked(base + 5);
125                *result.get_unchecked_mut(base + 6) =
126                    a.get_unchecked(base + 6) * b.get_unchecked(base + 6);
127                *result.get_unchecked_mut(base + 7) =
128                    a.get_unchecked(base + 7) * b.get_unchecked(base + 7);
129            }
130        }
131
132        // Handle remaining elements
133        let remainder_start = chunks * 8;
134        for i in 0..remainder {
135            unsafe {
136                *result.get_unchecked_mut(remainder_start + i) =
137                    a.get_unchecked(remainder_start + i) * b.get_unchecked(remainder_start + i);
138            }
139        }
140    }
141
142    /// Element-wise multiplication with optimization hints for f32 tensors (safe version)
143    pub fn mul_f32_optimized(a: &[f32], b: &[f32], result: &mut [f32]) -> Result<()> {
144        if a.len() != b.len() || a.len() != result.len() {
145            return Err(TensorError::ShapeMismatch {
146                operation: "SIMD mul_f32".to_string(),
147                expected: format!("arrays of length {}", a.len()),
148                got: format!("a: {}, b: {}, result: {}", a.len(), b.len(), result.len()),
149                context: Some(Box::new(ErrorContext::new())),
150            });
151        }
152
153        Self::mul_f32_unchecked(a, b, result);
154        Ok(())
155    }
156
157    /// Fast inline element-wise subtraction without bounds checking (for hot paths)
158    #[inline(always)]
159    pub fn sub_f32_unchecked(a: &[f32], b: &[f32], result: &mut [f32]) {
160        let len = a.len();
161        // For very small arrays, use simple vectorizable loop
162        if len < 32 {
163            for i in 0..len {
164                unsafe {
165                    *result.get_unchecked_mut(i) = a.get_unchecked(i) - b.get_unchecked(i);
166                }
167            }
168            return;
169        }
170        // For larger arrays, use optimal chunk size for SIMD (8 elements = 256 bits for AVX)
171        let chunks = len / 8;
172        let remainder = len % 8;
173        // Process 8 elements at a time - optimal for most SIMD instruction sets
174        for chunk in 0..chunks {
175            let base = chunk * 8;
176            unsafe {
177                *result.get_unchecked_mut(base) = a.get_unchecked(base) - b.get_unchecked(base);
178                *result.get_unchecked_mut(base + 1) =
179                    a.get_unchecked(base + 1) - b.get_unchecked(base + 1);
180                *result.get_unchecked_mut(base + 2) =
181                    a.get_unchecked(base + 2) - b.get_unchecked(base + 2);
182                *result.get_unchecked_mut(base + 3) =
183                    a.get_unchecked(base + 3) - b.get_unchecked(base + 3);
184                *result.get_unchecked_mut(base + 4) =
185                    a.get_unchecked(base + 4) - b.get_unchecked(base + 4);
186                *result.get_unchecked_mut(base + 5) =
187                    a.get_unchecked(base + 5) - b.get_unchecked(base + 5);
188                *result.get_unchecked_mut(base + 6) =
189                    a.get_unchecked(base + 6) - b.get_unchecked(base + 6);
190                *result.get_unchecked_mut(base + 7) =
191                    a.get_unchecked(base + 7) - b.get_unchecked(base + 7);
192            }
193        }
194        // Handle remaining elements
195        let remainder_start = chunks * 8;
196        for i in 0..remainder {
197            unsafe {
198                *result.get_unchecked_mut(remainder_start + i) =
199                    a.get_unchecked(remainder_start + i) - b.get_unchecked(remainder_start + i);
200            }
201        }
202    }
203
204    /// Element-wise subtraction with optimization hints for f32 tensors (safe version)
205    pub fn sub_f32_optimized(a: &[f32], b: &[f32], result: &mut [f32]) -> Result<()> {
206        if a.len() != b.len() || a.len() != result.len() {
207            return Err(TensorError::ShapeMismatch {
208                operation: "SIMD sub_f32".to_string(),
209                expected: format!("arrays of length {}", a.len()),
210                got: format!("a: {}, b: {}, result: {}", a.len(), b.len(), result.len()),
211                context: Some(Box::new(ErrorContext::new())),
212            });
213        }
214        Self::sub_f32_unchecked(a, b, result);
215        Ok(())
216    }
217
218    /// Ultra-high-performance fused multiply-add operation (a * b + c)
219    #[inline(always)]
220    pub fn fma_f32_unchecked(a: &[f32], b: &[f32], c: &[f32], result: &mut [f32]) {
221        let len = a.len();
222
223        // Process elements in chunks of 8 for optimal SIMD performance
224        let chunk_size = 8;
225        let full_chunks = len / chunk_size;
226
227        // Main processing loop with unrolling for SIMD optimization
228        for chunk in 0..full_chunks {
229            let base = chunk * chunk_size;
230
231            // Manually unrolled FMA operations for better vectorization
232            unsafe {
233                // Load and compute chunk
234                *result.get_unchecked_mut(base) =
235                    a.get_unchecked(base) * b.get_unchecked(base) + c.get_unchecked(base);
236                *result.get_unchecked_mut(base + 1) = a.get_unchecked(base + 1)
237                    * b.get_unchecked(base + 1)
238                    + c.get_unchecked(base + 1);
239                *result.get_unchecked_mut(base + 2) = a.get_unchecked(base + 2)
240                    * b.get_unchecked(base + 2)
241                    + c.get_unchecked(base + 2);
242                *result.get_unchecked_mut(base + 3) = a.get_unchecked(base + 3)
243                    * b.get_unchecked(base + 3)
244                    + c.get_unchecked(base + 3);
245                *result.get_unchecked_mut(base + 4) = a.get_unchecked(base + 4)
246                    * b.get_unchecked(base + 4)
247                    + c.get_unchecked(base + 4);
248                *result.get_unchecked_mut(base + 5) = a.get_unchecked(base + 5)
249                    * b.get_unchecked(base + 5)
250                    + c.get_unchecked(base + 5);
251                *result.get_unchecked_mut(base + 6) = a.get_unchecked(base + 6)
252                    * b.get_unchecked(base + 6)
253                    + c.get_unchecked(base + 6);
254                *result.get_unchecked_mut(base + 7) = a.get_unchecked(base + 7)
255                    * b.get_unchecked(base + 7)
256                    + c.get_unchecked(base + 7);
257            }
258        }
259
260        // Handle remaining elements
261        let remainder_start = full_chunks * chunk_size;
262        for i in remainder_start..len {
263            unsafe {
264                *result.get_unchecked_mut(i) =
265                    a.get_unchecked(i) * b.get_unchecked(i) + c.get_unchecked(i);
266            }
267        }
268    }
269
270    /// Batch vector operations with memory-friendly access patterns
271    pub fn batch_add_f32(batches: &[(&[f32], &[f32])], results: &mut [&mut [f32]]) -> Result<()> {
272        if batches.len() != results.len() {
273            return Err(TensorError::ShapeMismatch {
274                operation: "Batch SIMD add".to_string(),
275                expected: format!("{} result arrays", batches.len()),
276                got: format!("{} result arrays", results.len()),
277                context: Some(Box::new(ErrorContext::new())),
278            });
279        }
280
281        for (i, ((a, b), result)) in batches.iter().zip(results.iter_mut()).enumerate() {
282            if a.len() != b.len() || a.len() != result.len() {
283                return Err(TensorError::ShapeMismatch {
284                    operation: format!("Batch SIMD add (batch {})", i),
285                    expected: format!("arrays of length {}", a.len()),
286                    got: format!("a: {}, b: {}, result: {}", a.len(), b.len(), result.len()),
287                    context: Some(Box::new(ErrorContext::new())),
288                });
289            }
290
291            Self::add_f32_unchecked(a, b, result);
292        }
293
294        Ok(())
295    }
296
297    /// Auto-select the best implementation based on input characteristics
298    pub fn add_f32_auto(a: &[f32], b: &[f32], result: &mut [f32]) -> Result<()> {
299        let size = a.len();
300
301        if size < 64 {
302            // Use simple implementation for small arrays to avoid overhead
303            if a.len() != b.len() || a.len() != result.len() {
304                return Err(TensorError::ShapeMismatch {
305                    operation: "Simple add_f32".to_string(),
306                    expected: format!("arrays of length {}", a.len()),
307                    got: format!("a: {}, b: {}, result: {}", a.len(), b.len(), result.len()),
308                    context: Some(Box::new(ErrorContext::new())),
309                });
310            }
311
312            for i in 0..size {
313                result[i] = a[i] + b[i];
314            }
315        } else {
316            // Use optimized implementation for larger arrays
317            Self::add_f32_optimized(a, b, result)?;
318        }
319
320        Ok(())
321    }
322}
323
324#[cfg(test)]
325mod tests {
326    use super::*;
327    use approx::assert_relative_eq;
328
329    #[test]
330    fn test_add_f32_unchecked() {
331        let a = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0];
332        let b = vec![2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0];
333        let mut result = vec![0.0; 10];
334        let expected = vec![3.0, 5.0, 7.0, 9.0, 11.0, 13.0, 15.0, 17.0, 19.0, 21.0];
335
336        BasicOps::add_f32_unchecked(&a, &b, &mut result);
337
338        for (i, &val) in result.iter().enumerate() {
339            assert_relative_eq!(val, expected[i], epsilon = 1e-6);
340        }
341    }
342
343    #[test]
344    fn test_add_f32_optimized() {
345        let a = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0];
346        let b = vec![2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0];
347        let mut result = vec![0.0; 10];
348        let expected = vec![3.0, 5.0, 7.0, 9.0, 11.0, 13.0, 15.0, 17.0, 19.0, 21.0];
349
350        BasicOps::add_f32_optimized(&a, &b, &mut result)
351            .expect("test: add_f32_optimized should succeed");
352
353        for (i, &val) in result.iter().enumerate() {
354            assert_relative_eq!(val, expected[i], epsilon = 1e-6);
355        }
356    }
357
358    #[test]
359    fn test_mul_f32_optimized() {
360        let a = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
361        let b = vec![2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0];
362        let mut result = vec![0.0; 8];
363        let expected = [2.0, 6.0, 12.0, 20.0, 30.0, 42.0, 56.0, 72.0];
364
365        BasicOps::mul_f32_optimized(&a, &b, &mut result)
366            .expect("test: mul_f32_optimized should succeed");
367
368        for (i, &val) in result.iter().enumerate() {
369            assert_relative_eq!(val, expected[i], epsilon = 1e-6);
370        }
371    }
372
373    #[test]
374    fn test_fma_f32_unchecked() {
375        let a = vec![2.0, 3.0, 4.0, 5.0];
376        let b = vec![3.0, 4.0, 5.0, 6.0];
377        let c = vec![1.0, 2.0, 3.0, 4.0];
378        let mut result = vec![0.0; 4];
379        let expected = [7.0, 14.0, 23.0, 34.0]; // a*b+c
380
381        BasicOps::fma_f32_unchecked(&a, &b, &c, &mut result);
382
383        for (i, &val) in result.iter().enumerate() {
384            assert_relative_eq!(val, expected[i], epsilon = 1e-6);
385        }
386    }
387
388    #[test]
389    fn test_shape_mismatch_errors() {
390        let a = vec![1.0; 5];
391        let b = vec![1.0; 3]; // Wrong size
392        let mut result = vec![0.0; 5];
393
394        let error = BasicOps::add_f32_optimized(&a, &b, &mut result);
395        assert!(error.is_err());
396
397        let error = BasicOps::mul_f32_optimized(&a, &b, &mut result);
398        assert!(error.is_err());
399    }
400}
tenflowers_core/simd/basic_ops.rs

tenflowers_core/simd/
basic_ops.rs
