simd_kernels/kernels/arithmetic/
mod.rs

1// Copyright 2025 Peter Garfield Bower
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! # **Arithmetic Kernels Module** - *High-Performance Arithmetic*
16//!
17//! SIMD-optimised arithmetic operations for numeric arrays with null-aware semantics.
18//!
19//! ## Modules
20//! - **`dispatch`**: Smart dispatch layer selecting SIMD vs scalar implementations based on alignment
21//! - **`simd`**: SIMD-accelerated implementations using `std::simd` with portable vectorisation  
22//! - **`std`**: Scalar fallback implementations for compatibility and unaligned data
23//! - **`string`**: Specialised arithmetic operations for string concatenation and manipulation
24//!
25//! ## Operations
26//! Supports standard arithmetic operations (add, subtract, multiply, divide, remainder, power)
27//! plus fused multiply-add (FMA) for floating-point types with hardware acceleration.
28//!
29//! ## Scope
30//! **These do not leverage parallel-thread processing, as this is expected to be applied in the engine layer,
31//! which is app-specific.**.
32
33pub mod dispatch;
34#[cfg(feature = "simd")]
35pub mod simd;
36pub mod std;
37pub mod string;
38pub mod string_ops;
39
40// Shared tests for SIMD and Std
41
42#[cfg(test)]
43mod tests {
44    use crate::kernels::bitmask::merge_bitmasks_to_new;
45    use crate::structs::variants::float::FloatArray;
46    use crate::structs::variants::integer::IntegerArray;
47    use crate::{Bitmask, MaskedArray, vec64};
48
49    use crate::enums::operators::ArithmeticOperator;
50    use crate::kernels::arithmetic::dispatch::{
51        apply_float_f32, apply_float_f64, apply_fma_f32, apply_fma_f64, apply_int_i32,
52        apply_int_i64, apply_int_u32, apply_int_u64,
53    };
54    #[cfg(feature = "extended_numeric_types")]
55    use crate::kernels::arithmetic::dispatch::{
56        apply_int_i8, apply_int_i16, apply_int_u8, apply_int_u16,
57    };
58    #[cfg(feature = "simd")]
59    use crate::kernels::arithmetic::simd::int_dense_body_simd;
60
61    fn assert_int<T>(arr: &IntegerArray<T>, values: &[T], valid: Option<&[bool]>)
62    where
63        T: num_traits::PrimInt + std::fmt::Debug,
64    {
65        assert_eq!(arr.data.as_slice(), values);
66        match (valid, &arr.null_mask) {
67            (None, None) => {}
68            (Some(expected), Some(mask)) => {
69                for (i, bit) in expected.iter().enumerate() {
70                    assert_eq!(
71                        unsafe { mask.get_unchecked(i) },
72                        *bit,
73                        "mask mismatch at {i}"
74                    );
75                }
76            }
77            (None, Some(mask)) => {
78                assert!(mask.all_true(), "mask unexpectedly present");
79            }
80            (Some(_), None) => panic!("expected mask missing"),
81        }
82    }
83
84    fn assert_float<T>(arr: &FloatArray<T>, values: &[T], valid: Option<&[bool]>)
85    where
86        T: num_traits::Float + std::fmt::Debug,
87    {
88        assert_eq!(arr.data.as_slice(), values);
89        match (valid, &arr.null_mask) {
90            (None, None) => {}
91            (Some(expected), Some(mask)) => {
92                for (i, bit) in expected.iter().enumerate() {
93                    assert_eq!(
94                        unsafe { mask.get_unchecked(i) },
95                        *bit,
96                        "mask mismatch at {i}"
97                    );
98                }
99            }
100            (None, Some(mask)) => {
101                assert!(mask.all_true(), "mask unexpectedly present");
102            }
103            (Some(_), None) => panic!("expected mask missing"),
104        }
105    }
106
107    fn bitmask(bits: &[bool]) -> Bitmask {
108        let mut m = Bitmask::new_set_all(bits.len(), false);
109        for (i, b) in bits.iter().enumerate() {
110            unsafe { m.set_unchecked(i, *b) };
111        }
112        m
113    }
114
115    macro_rules! int_kernel_suite {
116        ($fn_dense:ident, $fn_masked:ident, $fn_empty:ident, $ty:ty, $apply_fn:ident) => {
117            #[test]
118            fn $fn_dense() {
119                let lhs = vec64![1, 4, 9, 16];
120                let rhs = vec64![1, 2, 3, 4];
121
122                let out = $apply_fn(&lhs, &rhs, ArithmeticOperator::Add, None).unwrap();
123                assert_int(
124                    &out,
125                    &IntegerArray::<$ty>::from_slice(&[2, 6, 12, 20]),
126                    None,
127                );
128
129                let out = $apply_fn(&lhs, &rhs, ArithmeticOperator::Subtract, None).unwrap();
130                assert_int(&out, &IntegerArray::<$ty>::from_slice(&[0, 2, 6, 12]), None);
131
132                let out = $apply_fn(&lhs, &rhs, ArithmeticOperator::Multiply, None).unwrap();
133                assert_int(
134                    &out,
135                    &IntegerArray::<$ty>::from_slice(&[1, 8, 27, 64]),
136                    None,
137                );
138
139                let out = $apply_fn(&lhs, &rhs, ArithmeticOperator::Divide, None).unwrap();
140                assert_int(&out, &IntegerArray::<$ty>::from_slice(&[1, 2, 3, 4]), None);
141
142                let out = $apply_fn(&lhs, &rhs, ArithmeticOperator::Remainder, None).unwrap();
143                assert_int(&out, &IntegerArray::<$ty>::from_slice(&[0, 0, 0, 0]), None);
144
145                let expected: Vec<$ty> = lhs
146                    .iter()
147                    .zip(rhs.iter())
148                    .map(|(&a, &b)| {
149                        let mut acc = <$ty as num_traits::One>::one();
150                        for _ in 0..(b as u32) {
151                            acc = acc.wrapping_mul(a);
152                        }
153                        acc
154                    })
155                    .collect();
156                let out = $apply_fn(&lhs, &rhs, ArithmeticOperator::Power, None).unwrap();
157                assert_int(&out, &IntegerArray::<$ty>::from_slice(&expected), None);
158
159                // Division by zero should panic
160                let rhs_divzero: &[$ty] = &[0, 0, 0, 0];
161                let result = std::panic::catch_unwind(|| {
162                    $apply_fn(&lhs, rhs_divzero, ArithmeticOperator::Divide, None).unwrap()
163                });
164                assert!(
165                    result.is_err(),
166                    "Dense integer kernel division by zero must panic"
167                );
168
169                let result = std::panic::catch_unwind(|| {
170                    $apply_fn(&lhs, rhs_divzero, ArithmeticOperator::Remainder, None).unwrap()
171                });
172                assert!(
173                    result.is_err(),
174                    "Dense integer kernel remainder by zero must panic"
175                );
176            }
177
178            #[test]
179            fn $fn_masked() {
180                let lhs = vec64![10, 20, 30, 40];
181                let rhs = vec64![2, 0, 3, 5];
182                let mask = bitmask(&[true, false, true, false]);
183
184                // Division: mask==true and rhs!=0 are valid, otherwise null
185                let out = $apply_fn(&lhs, &rhs, ArithmeticOperator::Divide, Some(&mask)).unwrap();
186                assert_int(
187                    &out,
188                    &IntegerArray::<$ty>::from_slice(&[5, 0, 10, 0]),
189                    Some(&[true, false, true, false]),
190                );
191
192                // Remainder with mask, matching above
193                let out =
194                    $apply_fn(&lhs, &rhs, ArithmeticOperator::Remainder, Some(&mask)).unwrap();
195                assert_int(
196                    &out,
197                    &IntegerArray::<$ty>::from_slice(&[0, 0, 0, 0]),
198                    Some(&[true, false, true, false]),
199                );
200
201                // Division by zero where mask is true but rhs is zero must yield null in mask (false) and output 0
202                let mask_divzero = bitmask(&[true, true, true, true]);
203                let rhs_divzero: &[$ty] = &[1, 0, 2, 0];
204                let lhs2: &[$ty] = &[100, 100, 100, 100];
205
206                let out = $apply_fn(
207                    lhs2,
208                    rhs_divzero,
209                    ArithmeticOperator::Divide,
210                    Some(&mask_divzero),
211                )
212                .unwrap();
213                assert_int(
214                    &out,
215                    &IntegerArray::<$ty>::from_slice(&[100, 0, 50, 0]),
216                    Some(&[true, false, true, false]),
217                );
218            }
219
220            #[test]
221            fn $fn_empty() {
222                let lhs = vec64![];
223                let rhs = vec64![];
224                let out = $apply_fn(&lhs, &rhs, ArithmeticOperator::Add, None).unwrap();
225                assert!(out.is_empty());
226            }
227        };
228    }
229
230    #[cfg(feature = "extended_numeric_types")]
231    int_kernel_suite!(
232        apply_int_i8_dense,
233        apply_int_i8_masked,
234        apply_int_i8_empty,
235        i8,
236        apply_int_i8
237    );
238    #[cfg(feature = "extended_numeric_types")]
239    int_kernel_suite!(
240        apply_int_u8_dense,
241        apply_int_u8_masked,
242        apply_int_u8_empty,
243        u8,
244        apply_int_u8
245    );
246    #[cfg(feature = "extended_numeric_types")]
247    int_kernel_suite!(
248        apply_int_i16_dense,
249        apply_int_i16_masked,
250        apply_int_i16_empty,
251        i16,
252        apply_int_i16
253    );
254    #[cfg(feature = "extended_numeric_types")]
255    int_kernel_suite!(
256        apply_int_u16_dense,
257        apply_int_u16_masked,
258        apply_int_u16_empty,
259        u16,
260        apply_int_u16
261    );
262    int_kernel_suite!(
263        apply_int_i32_dense,
264        apply_int_i32_masked,
265        apply_int_i32_empty,
266        i32,
267        apply_int_i32
268    );
269    int_kernel_suite!(
270        apply_int_u32_dense,
271        apply_int_u32_masked,
272        apply_int_u32_empty,
273        u32,
274        apply_int_u32
275    );
276    int_kernel_suite!(
277        apply_int_i64_dense,
278        apply_int_i64_masked,
279        apply_int_i64_empty,
280        i64,
281        apply_int_i64
282    );
283    int_kernel_suite!(
284        apply_int_u64_dense,
285        apply_int_u64_masked,
286        apply_int_u64_empty,
287        u64,
288        apply_int_u64
289    );
290
291    macro_rules! float_kernel_suite {
292        ($test_fn:ident, $ty:ty, $apply_fn:ident, $eps:expr) => {
293            #[test]
294            fn $test_fn() {
295                let lhs = vec64![1.0, 4.0, 9.0, 16.0];
296                let rhs = vec64![0.5, 2.0, 3.0, 4.0];
297
298                let lhs: &[$ty] = lhs.as_slice();
299                let rhs: &[$ty] = rhs.as_slice();
300
301                let arr = $apply_fn(lhs, rhs, ArithmeticOperator::Add, None).unwrap();
302                assert_eq!(arr.data.as_slice(), &[1.5 as $ty, 6.0, 12.0, 20.0]);
303
304                let arr = $apply_fn(lhs, rhs, ArithmeticOperator::Subtract, None).unwrap();
305                assert_eq!(arr.data.as_slice(), &[0.5 as $ty, 2.0, 6.0, 12.0]);
306
307                let arr = $apply_fn(lhs, rhs, ArithmeticOperator::Multiply, None).unwrap();
308                assert_eq!(arr.data.as_slice(), &[0.5 as $ty, 8.0, 27.0, 64.0]);
309
310                let arr = $apply_fn(lhs, rhs, ArithmeticOperator::Divide, None).unwrap();
311                assert_eq!(arr.data.as_slice(), &[2.0 as $ty, 2.0, 3.0, 4.0]);
312
313                let arr = $apply_fn(lhs, rhs, ArithmeticOperator::Remainder, None).unwrap();
314                assert!(
315                    arr.data
316                        .as_slice()
317                        .iter()
318                        .zip(
319                            [1.0 % 0.5, 4.0 % 2.0, 9.0 % 3.0, 16.0 % 4.0]
320                                .iter()
321                                .map(|&x| x as $ty)
322                        )
323                        .all(|(a, b)| (*a - b).abs() < $eps)
324                );
325
326                let arr = $apply_fn(lhs, rhs, ArithmeticOperator::Power, None).unwrap();
327                let expected: Vec<$ty> = lhs
328                    .iter()
329                    .zip(rhs.iter())
330                    .map(|(&a, &b)| (b * a.ln()).exp())
331                    .collect();
332                assert!(
333                    arr.data
334                        .as_slice()
335                        .iter()
336                        .zip(expected.iter())
337                        .all(|(a, b)| (*a - *b).abs() < $eps)
338                );
339
340                // Division by zero for floats yields Inf/NaN, never panics
341                let rhs_divzero: &[$ty] = &[0.0, 0.0, 0.0, 0.0];
342                let arr = $apply_fn(lhs, rhs_divzero, ArithmeticOperator::Divide, None).unwrap();
343                assert!(
344                    arr.data.iter().all(|&x| x.is_infinite()),
345                    "Float division by zero should yield Inf"
346                );
347
348                let arr = $apply_fn(lhs, rhs_divzero, ArithmeticOperator::Remainder, None).unwrap();
349                assert!(
350                    arr.data.iter().all(|&x| x.is_nan()),
351                    "Float remainder by zero should yield NaN"
352                );
353
354                // Masked test
355                let mask = bitmask(&[true, false, true, false]);
356                let arr = $apply_fn(lhs, rhs, ArithmeticOperator::Multiply, Some(&mask)).unwrap();
357                assert_eq!(arr.data.as_slice(), &[0.5 as $ty, 0.0, 27.0, 0.0]);
358                assert_eq!(arr.null_mask.as_ref().unwrap().len(), 4);
359
360                // Empty
361                let arr = $apply_fn(&[], &[], ArithmeticOperator::Add, None).unwrap();
362                assert!(arr.is_empty());
363            }
364        };
365    }
366
367    float_kernel_suite!(apply_float_f32_dense, f32, apply_float_f32, 1e-6f32);
368    float_kernel_suite!(apply_float_f64_dense, f64, apply_float_f64, 1e-12f64);
369
370    #[test]
371    fn fma_f32() {
372        let lhs = vec64![1.0f32, 2.0, 3.0];
373        let rhs = vec64![4.0f32, 5.0, 6.0];
374        let acc = vec64![0.5f32, 0.5, 0.5];
375        let out = apply_fma_f32(&lhs, &rhs, &acc, None).unwrap();
376        assert_float(&out, &[4.5, 10.5, 18.5], None);
377
378        let mask = bitmask(&[true, false, true]);
379        let out = apply_fma_f32(&lhs, &rhs, &acc, Some(&mask)).unwrap();
380        assert_float(&out, &[4.5, 0.0, 18.5], Some(&[true, false, true]));
381
382        let out = apply_fma_f32(&[], &[], &[], None).unwrap();
383        assert!(out.is_empty());
384    }
385
386    #[test]
387    fn fma_f64() {
388        let lhs = vec64![1.0f64, 2.0, 3.0];
389        let rhs = vec64![4.0f64, 5.0, 6.0];
390        let acc = vec64![0.5f64, 0.5, 0.5];
391        let out = apply_fma_f64(&lhs, &rhs, &acc, None).unwrap();
392        assert_float(&out, &[4.5, 10.5, 18.5], None);
393
394        let mask = bitmask(&[true, false, true]);
395        let out = apply_fma_f64(&lhs, &rhs, &acc, Some(&mask)).unwrap();
396        assert_float(&out, &[4.5, 0.0, 18.5], Some(&[true, false, true]));
397    }
398
399    #[test]
400    fn merge_masks_correctness() {
401        let a = bitmask(&[true, false, true, true]);
402        let b = bitmask(&[true, true, false, true]);
403        let merged = merge_bitmasks_to_new(Some(&a), Some(&b), 4).unwrap();
404        let expected = vec![true, false, false, true];
405        let merged_vec: Vec<bool> = (0..4).map(|i| merged.get(i)).collect();
406        assert_eq!(merged_vec, expected);
407    }
408
409    // Datetime Kernels
410    #[cfg(feature = "datetime")]
411    use crate::structs::variants::datetime::DatetimeArray;
412
413    #[cfg(feature = "datetime")]
414    use crate::kernels::arithmetic::dispatch::apply_datetime_i64;
415
416    #[cfg(feature = "datetime")]
417    #[test]
418    fn datetime_add() {
419        let lhs = DatetimeArray::<i64>::from_slice(&[1_000i64, 2_000, 3_000], None);
420        let rhs = DatetimeArray::<i64>::from_slice(&[10, 20, 30], None);
421        let lhs_slice = (&lhs, 0, lhs.len());
422        let rhs_slice = (&rhs, 0, rhs.len());
423        let out = apply_datetime_i64(lhs_slice, rhs_slice, ArithmeticOperator::Add).unwrap();
424        assert_eq!(out.data.as_slice(), &[1_010, 2_020, 3_030]);
425        assert!(out.null_mask.is_none());
426    }
427
428    #[cfg(feature = "datetime")]
429    #[test]
430    fn datetime_all_ops() {
431        let lhs = DatetimeArray::<i64>::from_slice(&[10, 20, 30, 40], None);
432        let rhs = DatetimeArray::<i64>::from_slice(&[1, 2, 3, 4], None);
433        let lhs_slice = (&lhs, 0, lhs.len());
434        let rhs_slice = (&rhs, 0, rhs.len());
435
436        let out = apply_datetime_i64(lhs_slice, rhs_slice, ArithmeticOperator::Add).unwrap();
437        assert_eq!(out.data.as_slice(), &[11, 22, 33, 44]);
438
439        let out = apply_datetime_i64(lhs_slice, rhs_slice, ArithmeticOperator::Subtract).unwrap();
440        assert_eq!(out.data.as_slice(), &[9, 18, 27, 36]);
441
442        let out = apply_datetime_i64(lhs_slice, rhs_slice, ArithmeticOperator::Multiply).unwrap();
443        assert_eq!(out.data.as_slice(), &[10, 40, 90, 160]);
444
445        let out = apply_datetime_i64(lhs_slice, rhs_slice, ArithmeticOperator::Divide).unwrap();
446        assert_eq!(out.data.as_slice(), &[10, 10, 10, 10]);
447
448        let out = apply_datetime_i64(lhs_slice, rhs_slice, ArithmeticOperator::Remainder).unwrap();
449        assert_eq!(out.data.as_slice(), &[0, 0, 0, 0]);
450
451        let out = apply_datetime_i64(lhs_slice, rhs_slice, ArithmeticOperator::Power).unwrap();
452        assert_eq!(
453            out.data.as_slice(),
454            &[10_i64.pow(1), 20_i64.pow(2), 30_i64.pow(3), 40_i64.pow(4)]
455        );
456    }
457
458    #[cfg(feature = "datetime")]
459    #[test]
460    fn datetime_masked_and_empty() {
461        let lhs = DatetimeArray::<i64>::from_slice(&[10, 20, 30, 40], None);
462        let rhs = DatetimeArray::<i64>::from_slice(&[1, 2, 3, 4], None);
463        let mask = bitmask(&[true, false, true, true]);
464        let lhs_slice = (&lhs, 0, lhs.len());
465        let rhs_slice = (&rhs, 0, rhs.len());
466
467        let out = apply_datetime_i64(lhs_slice, rhs_slice, ArithmeticOperator::Add).unwrap();
468        assert_eq!(out.data.as_slice(), &[11, 22, 33, 44]);
469
470        // Masked
471        let mut lhs_masked = lhs.clone();
472        lhs_masked.null_mask = Some(mask.clone());
473        let lhs_slice_masked = (&lhs_masked, 0, lhs_masked.len());
474        let out = apply_datetime_i64(lhs_slice_masked, rhs_slice, ArithmeticOperator::Add).unwrap();
475        let expected = vec![11, 0, 33, 44];
476        let mask_vec: Vec<bool> = (0..4).map(|i| mask.get(i)).collect();
477        assert_eq!(out.data.as_slice(), &expected);
478        assert_eq!(
479            out.null_mask
480                .as_ref()
481                .map(|m| (0..4).map(|i| m.get(i)).collect::<Vec<_>>()),
482            Some(mask_vec)
483        );
484
485        // Empty
486        let lhs_empty = DatetimeArray::<i64>::from_slice(&[], None);
487        let rhs_empty = DatetimeArray::<i64>::from_slice(&[], None);
488        let lhs_slice = (&lhs_empty, 0, lhs_empty.len());
489        let rhs_slice = (&rhs_empty, 0, rhs_empty.len());
490        let out = apply_datetime_i64(lhs_slice, rhs_slice, ArithmeticOperator::Add).unwrap();
491        assert!(out.is_empty());
492    }
493
494    #[cfg(feature = "datetime")]
495    #[test]
496    #[should_panic(expected = "apply_datetime: length mismatch")]
497    fn datetime_len_mismatch_panics() {
498        let lhs = DatetimeArray::<i64>::from_slice(&[1_000i64, 2_000], None);
499        let rhs = DatetimeArray::<i64>::from_slice(&[10], None);
500        let lhs_slice = (&lhs, 0, lhs.len());
501        let rhs_slice = (&rhs, 0, rhs.len());
502        let _ = apply_datetime_i64(lhs_slice, rhs_slice, ArithmeticOperator::Add).unwrap();
503    }
504
505    #[cfg(feature = "simd")]
506    #[test]
507    fn test_int_dense_power_short_vs_long_input_simd() {
508        let lhs_short = vec64![2u32; 16];
509        let rhs_short = vec64![10u32; 16];
510        let mut out_short = vec64![0u32; 16];
511
512        let lhs_long = vec64![2u32; 128];
513        let rhs_long = vec64![10u32; 128];
514        let mut out_long = vec64![0u32; 128];
515
516        int_dense_body_simd::<u32, 4>(
517            ArithmeticOperator::Power,
518            &lhs_short,
519            &rhs_short,
520            &mut out_short,
521        );
522        int_dense_body_simd::<u32, 4>(
523            ArithmeticOperator::Power,
524            &lhs_long,
525            &rhs_long,
526            &mut out_long,
527        );
528
529        for &v in out_short.iter() {
530            assert_eq!(v, 1024);
531        }
532        for &v in out_long.iter() {
533            assert_eq!(v, 1024);
534        }
535    }
536}
simd_kernels/kernels/arithmetic/mod.rs

simd_kernels/kernels/arithmetic/
mod.rs
