simd_kernels/kernels/arithmetic/
simd.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//! # **SIMD Arithmetic Kernels Module** - *High-Performance Arithmetic*
16//!
17//! Inner SIMD-accelerated implementations using `std::simd` for maximum performance on modern hardware.
18//! Prefer dispatch.rs for easily handling the general case, otherwise you can use these inner functions
19//! directly (e.g., "dense_simd") vs. "maybe masked, maybe simd".
20//!
21//! ## Overview
22//! - **Portable SIMD**: Uses `std::simd` for cross-platform vectorisation with compile-time lane optimisation
23//! - **Null masks**: Dense (no nulls) and masked variants for Arrow-compatible null handling.
24//!   These are uniified in dispatch.rs, and opting out of masking yields no performance penalty.
25//! - **Type support**: Integer and floating-point arithmetic with specialised FMA operations
26//! - **Safety**: All unsafe operations are bounds-checked or guaranteed by caller invariants
27//!
28//! ## Architecture Notes
29//! - Building blocks for higher-level dispatch layers, or for low-level hot loops
30//! where one wants to fully avoid abstraction overhead.
31//! - Parallelisation intentionally excluded to allow flexible chunking strategies
32//! - Power operations fall back to scalar for integers, use logarithmic computation for floats
33
34include!(concat!(env!("OUT_DIR"), "/simd_lanes.rs"));
35
36use core::simd::{Mask, Simd, SimdElement};
37use std::ops::{Add, Div, Mul, Rem, Sub};
38use std::simd::cmp::SimdPartialEq;
39use std::simd::{Select, StdFloat};
40
41use crate::Bitmask;
42use num_traits::{One, PrimInt, ToPrimitive, WrappingAdd, WrappingMul, WrappingSub, Zero};
43
44use crate::enums::operators::ArithmeticOperator;
45use crate::kernels::bitmask::simd::all_true_mask_simd;
46use crate::utils::{simd_mask, write_simd_mask_bits};
47
48/// SIMD integer arithmetic kernel for dense arrays (no nulls).
49/// Vectorised operations with scalar fallback for power operations and array tails.
50/// Panics on division/remainder by zero (consistent with scalar behaviour).
51#[inline(always)]
52pub fn int_dense_body_simd<T, const LANES: usize>(
53    op: ArithmeticOperator,
54    lhs: &[T],
55    rhs: &[T],
56    out: &mut [T],
57) where
58    T: Copy + One + PrimInt + ToPrimitive + Zero + SimdElement + WrappingMul,
59    Simd<T, LANES>: Add<Output = Simd<T, LANES>>
60        + Sub<Output = Simd<T, LANES>>
61        + Mul<Output = Simd<T, LANES>>
62        + Div<Output = Simd<T, LANES>>
63        + Rem<Output = Simd<T, LANES>>,
64{
65    let n = lhs.len();
66    let mut vectorisable = n / LANES * LANES;
67    let mut i = 0;
68    while i < vectorisable {
69        let a = Simd::<T, LANES>::from_slice(&lhs[i..i + LANES]);
70        let b = Simd::<T, LANES>::from_slice(&rhs[i..i + LANES]);
71        let r = match op {
72            ArithmeticOperator::Add => a + b,
73            ArithmeticOperator::Subtract => a - b,
74            ArithmeticOperator::Multiply => a * b,
75            ArithmeticOperator::Divide => a / b, // Panics if divisor is zero
76            ArithmeticOperator::Remainder => a % b, // Panics if divisor is zero
77            ArithmeticOperator::Power | ArithmeticOperator::FloorDiv => {
78                vectorisable = 0;
79                break;
80            }
81        };
82        r.copy_to_slice(&mut out[i..i + LANES]);
83        i += LANES;
84    }
85
86    // Scalar tail
87    for idx in vectorisable..n {
88        out[idx] = match op {
89            ArithmeticOperator::Add => lhs[idx] + rhs[idx],
90            ArithmeticOperator::Subtract => lhs[idx] - rhs[idx],
91            ArithmeticOperator::Multiply => lhs[idx] * rhs[idx],
92            ArithmeticOperator::Divide => lhs[idx] / rhs[idx], // Panics if divisor is zero
93            ArithmeticOperator::Remainder => lhs[idx] % rhs[idx], // Panics if divisor is zero
94            ArithmeticOperator::Power => {
95                let mut acc = T::one();
96                let exp = rhs[idx].to_u32().unwrap_or(0);
97                for _ in 0..exp {
98                    acc = acc.wrapping_mul(&lhs[idx]);
99                }
100                acc
101            }
102            ArithmeticOperator::FloorDiv => {
103                if rhs[idx] == T::zero() {
104                    panic!("Floor division by zero")
105                } else {
106                    let d = lhs[idx] / rhs[idx];
107                    let r = lhs[idx] % rhs[idx];
108                    if r != T::zero() && (lhs[idx] ^ rhs[idx]) < T::zero() { d - T::one() } else { d }
109                }
110            }
111        };
112    }
113}
114
115/// SIMD integer arithmetic kernel with null mask support.
116/// Division/remainder by zero produces null results (mask=false) rather than panicking.
117#[inline(always)]
118pub fn int_masked_body_simd<T, const LANES: usize>(
119    op: ArithmeticOperator,
120    lhs: &[T],
121    rhs: &[T],
122    mask: &Bitmask,
123    out: &mut [T],
124    out_mask: &mut Bitmask,
125) where
126    T: Copy
127        + PrimInt
128        + ToPrimitive
129        + Zero
130        + One
131        + SimdElement
132        + PartialEq
133        + WrappingAdd
134        + WrappingMul
135        + WrappingSub,
136    Simd<T, LANES>: Add<Output = Simd<T, LANES>>
137        + SimdPartialEq<Mask = Mask<<T as SimdElement>::Mask, LANES>>
138        + Sub<Output = Simd<T, LANES>>
139        + Mul<Output = Simd<T, LANES>>
140        + Div<Output = Simd<T, LANES>>
141        + Rem<Output = Simd<T, LANES>>,
142{
143    let n = lhs.len();
144    let dense = all_true_mask_simd::<LANES>(mask);
145
146    /* If dense, we unfortunately need to near-replicate the dense implementation
147    as that dedicated function panics on `div/0` as it needs to stay mask-free,
148    to support varied workloads. This one works on the same dense principles,
149    but substitutes the null mask when any div/0 issues occur. */
150
151    if dense {
152        // This block replaces the int_dense_body_simd call and handles masking for div/rem
153        let vectorisable = n / LANES * LANES;
154        let mut i = 0;
155        while i < vectorisable {
156            let a = Simd::<T, LANES>::from_slice(&lhs[i..i + LANES]);
157            let b = Simd::<T, LANES>::from_slice(&rhs[i..i + LANES]);
158
159            let (r, valid): (Simd<T, LANES>, Mask<<T as SimdElement>::Mask, LANES>) = match op {
160                ArithmeticOperator::Add => (a + b, Mask::splat(true)),
161                ArithmeticOperator::Subtract => (a - b, Mask::splat(true)),
162                ArithmeticOperator::Multiply => (a * b, Mask::splat(true)),
163                ArithmeticOperator::Power => {
164                    let mut tmp = [T::zero(); LANES];
165                    for l in 0..LANES {
166                        tmp[l] = a[l].pow(b[l].to_u32().unwrap_or(0));
167                    }
168                    (Simd::<T, LANES>::from_array(tmp), Mask::splat(true))
169                }
170                ArithmeticOperator::Divide | ArithmeticOperator::Remainder => {
171                    let div_zero = b.simd_eq(Simd::splat(T::zero()));
172                    let valid = !div_zero;
173                    let safe_b = div_zero.select(Simd::splat(T::one()), b);
174                    let r = match op {
175                        ArithmeticOperator::Divide => a / safe_b,
176                        ArithmeticOperator::Remainder => a % safe_b,
177                        _ => unreachable!(),
178                    };
179                    let r = div_zero.select(Simd::splat(T::zero()), r);
180                    (r, valid)
181                }
182                ArithmeticOperator::FloorDiv => {
183                    let div_zero = b.simd_eq(Simd::splat(T::zero()));
184                    let valid = !div_zero;
185                    // Per-lane floor division with sign correction
186                    let mut tmp = [T::zero(); LANES];
187                    for l in 0..LANES {
188                        if b[l] == T::zero() {
189                            tmp[l] = T::zero();
190                        } else {
191                            let d = a[l] / b[l];
192                            let r = a[l] % b[l];
193                            tmp[l] = if r != T::zero() && (a[l] ^ b[l]) < T::zero() { d - T::one() } else { d };
194                        }
195                    }
196                    (Simd::<T, LANES>::from_array(tmp), valid)
197                }
198            };
199            r.copy_to_slice(&mut out[i..i + LANES]);
200            // Write the out_mask based on the op
201            write_simd_mask_bits(out_mask, i, valid);
202            i += LANES;
203        }
204        // Scalar tail
205        for idx in vectorisable..n {
206            match op {
207                ArithmeticOperator::Add => {
208                    out[idx] = lhs[idx].wrapping_add(&rhs[idx]);
209                    unsafe {
210                        out_mask.set_unchecked(idx, true);
211                    }
212                }
213                ArithmeticOperator::Subtract => {
214                    out[idx] = lhs[idx].wrapping_sub(&rhs[idx]);
215                    unsafe {
216                        out_mask.set_unchecked(idx, true);
217                    }
218                }
219                ArithmeticOperator::Multiply => {
220                    out[idx] = lhs[idx].wrapping_mul(&rhs[idx]);
221                    unsafe {
222                        out_mask.set_unchecked(idx, true);
223                    }
224                }
225                ArithmeticOperator::Power => {
226                    out[idx] = lhs[idx].pow(rhs[idx].to_u32().unwrap_or(0));
227                    unsafe {
228                        out_mask.set_unchecked(idx, true);
229                    }
230                }
231                ArithmeticOperator::Divide | ArithmeticOperator::Remainder => {
232                    if rhs[idx] == T::zero() {
233                        out[idx] = T::zero();
234                        unsafe {
235                            out_mask.set_unchecked(idx, false);
236                        }
237                    } else {
238                        out[idx] = match op {
239                            ArithmeticOperator::Divide => lhs[idx] / rhs[idx],
240                            ArithmeticOperator::Remainder => lhs[idx] % rhs[idx],
241                            _ => unreachable!(),
242                        };
243                        unsafe {
244                            out_mask.set_unchecked(idx, true);
245                        }
246                    }
247                }
248                ArithmeticOperator::FloorDiv => {
249                    if rhs[idx] == T::zero() {
250                        out[idx] = T::zero();
251                        unsafe {
252                            out_mask.set_unchecked(idx, false);
253                        }
254                    } else {
255                        let d = lhs[idx] / rhs[idx];
256                        let r = lhs[idx] % rhs[idx];
257                        out[idx] = if r != T::zero() && (lhs[idx] ^ rhs[idx]) < T::zero() { d - T::one() } else { d };
258                        unsafe {
259                            out_mask.set_unchecked(idx, true);
260                        }
261                    }
262                }
263            }
264        }
265        return;
266    }
267
268    let mut i = 0;
269    while i + LANES <= n {
270        let a = Simd::<T, LANES>::from_slice(&lhs[i..i + LANES]);
271        let b = Simd::<T, LANES>::from_slice(&rhs[i..i + LANES]);
272        let m_src: Mask<_, LANES> = simd_mask::<_, LANES>(mask, i, n); // validity mask
273
274        // divisor-is-zero mask
275        let div_zero: Mask<_, LANES> = b.simd_eq(Simd::splat(T::zero()));
276
277        // ── compute result ───────────────────────────────────────────────────
278        let res = match op {
279            ArithmeticOperator::Add => a + b,
280            ArithmeticOperator::Subtract => a - b,
281            ArithmeticOperator::Multiply => a * b,
282            ArithmeticOperator::Divide => {
283                let safe_b = div_zero.select(Simd::splat(T::one()), b); // 0 -> 1
284                let q = a / safe_b;
285                div_zero.select(Simd::splat(T::zero()), q) // restore 0
286            }
287            ArithmeticOperator::Remainder => {
288                let safe_b = div_zero.select(Simd::splat(T::one()), b);
289                let r = a % safe_b;
290                div_zero.select(Simd::splat(T::zero()), r)
291            }
292            ArithmeticOperator::Power => {
293                // scalar per-lane power
294                let mut tmp = [T::zero(); LANES];
295                for l in 0..LANES {
296                    tmp[l] = a[l].pow(b[l].to_u32().unwrap_or(0));
297                }
298                Simd::<T, LANES>::from_array(tmp)
299            }
300            ArithmeticOperator::FloorDiv => {
301                // Per-lane floor division with sign correction
302                let mut tmp = [T::zero(); LANES];
303                for l in 0..LANES {
304                    if b[l] != T::zero() {
305                        let d = a[l] / b[l];
306                        let r = a[l] % b[l];
307                        tmp[l] = if r != T::zero() && (a[l] ^ b[l]) < T::zero() { d - T::one() } else { d };
308                    }
309                }
310                Simd::<T, LANES>::from_array(tmp)
311            }
312        };
313
314        // apply source validity mask, write results
315        let selected = m_src.select(res, Simd::splat(T::zero()));
316        selected.copy_to_slice(&mut out[i..i + LANES]);
317
318        // write out-mask bits: combine source mask with div-by-zero validity
319        let final_mask = match op {
320            ArithmeticOperator::Divide | ArithmeticOperator::Remainder | ArithmeticOperator::FloorDiv => {
321                // Valid iff source is valid and not dividing by zero
322                m_src & !div_zero
323            }
324            _ => m_src,
325        };
326        write_simd_mask_bits(out_mask, i, final_mask);
327        i += LANES;
328    }
329
330    // scalar tail
331    for j in i..n {
332        let valid = unsafe { mask.get_unchecked(j) };
333        if valid {
334            let (result, final_valid) = match op {
335                ArithmeticOperator::Add => (lhs[j].wrapping_add(&rhs[j]), true),
336                ArithmeticOperator::Subtract => (lhs[j].wrapping_sub(&rhs[j]), true),
337                ArithmeticOperator::Multiply => (lhs[j].wrapping_mul(&rhs[j]), true),
338                ArithmeticOperator::Divide => {
339                    if rhs[j] == T::zero() {
340                        (T::zero(), false) // division by zero -> invalid
341                    } else {
342                        (lhs[j] / rhs[j], true)
343                    }
344                }
345                ArithmeticOperator::Remainder => {
346                    if rhs[j] == T::zero() {
347                        (T::zero(), false) // remainder by zero -> invalid
348                    } else {
349                        (lhs[j] % rhs[j], true)
350                    }
351                }
352                ArithmeticOperator::Power => (lhs[j].pow(rhs[j].to_u32().unwrap_or(0)), true),
353                ArithmeticOperator::FloorDiv => {
354                    if rhs[j] == T::zero() {
355                        (T::zero(), false)
356                    } else {
357                        let d = lhs[j] / rhs[j];
358                        let r = lhs[j] % rhs[j];
359                        if r != T::zero() && (lhs[j] ^ rhs[j]) < T::zero() { (d - T::one(), true) } else { (d, true) }
360                    }
361                }
362            };
363            out[j] = result;
364            unsafe { out_mask.set_unchecked(j, final_valid) };
365        } else {
366            out[j] = T::zero();
367            unsafe { out_mask.set_unchecked(j, false) };
368        }
369    }
370}
371
372/// SIMD f32 arithmetic kernel with null mask support.
373/// Preserves IEEE 754 semantics: division by zero produces Inf/NaN, no exceptions.
374/// Power operations use scalar fallback with logarithmic computation.
375#[inline(always)]
376pub fn float_masked_body_f32_simd<const LANES: usize>(
377    op: ArithmeticOperator,
378    lhs: &[f32],
379    rhs: &[f32],
380    mask: &Bitmask,
381    out: &mut [f32],
382    out_mask: &mut Bitmask,
383) {
384    type M = <f32 as SimdElement>::Mask;
385
386    let n = lhs.len();
387    let dense = all_true_mask_simd::<LANES>(mask);
388
389    if dense {
390        float_dense_body_f32_simd::<LANES>(op, lhs, rhs, out);
391        out_mask.fill(true);
392        return;
393    }
394
395    let mut i = 0;
396    while i + LANES <= n {
397        let a = Simd::<f32, LANES>::from_slice(&lhs[i..i + LANES]);
398        let b = Simd::<f32, LANES>::from_slice(&rhs[i..i + LANES]);
399        let m: Mask<M, LANES> = simd_mask::<M, LANES>(mask, i, n);
400
401        let res = match op {
402            ArithmeticOperator::Add => a + b,
403            ArithmeticOperator::Subtract => a - b,
404            ArithmeticOperator::Multiply => a * b,
405            ArithmeticOperator::Divide => a / b,
406            ArithmeticOperator::Remainder => a % b,
407            ArithmeticOperator::Power => (b * a.ln()).exp(),
408            ArithmeticOperator::FloorDiv => (a / b).floor(),
409        };
410
411        let selected = m.select(res, Simd::<f32, LANES>::splat(0.0));
412        selected.copy_to_slice(&mut out[i..i + LANES]);
413
414        write_simd_mask_bits(out_mask, i, m);
415        i += LANES;
416    }
417
418    // Tail often caused by `n % LANES != 0`; uses scalar fallback
419    for j in i..n {
420        let valid = unsafe { mask.get_unchecked(j) };
421        if valid {
422            out[j] = match op {
423                ArithmeticOperator::Add => lhs[j] + rhs[j],
424                ArithmeticOperator::Subtract => lhs[j] - rhs[j],
425                ArithmeticOperator::Multiply => lhs[j] * rhs[j],
426                ArithmeticOperator::Divide => lhs[j] / rhs[j],
427                ArithmeticOperator::Remainder => lhs[j] % rhs[j],
428                ArithmeticOperator::Power => (rhs[j] * lhs[j].ln()).exp(),
429                ArithmeticOperator::FloorDiv => (lhs[j] / rhs[j]).floor(),
430            };
431            unsafe { out_mask.set_unchecked(j, true) };
432        } else {
433            out[j] = 0.0;
434            unsafe { out_mask.set_unchecked(j, false) };
435        }
436    }
437}
438
439/// SIMD f64 arithmetic kernel with null mask support.
440/// Preserves IEEE 754 semantics: division by zero produces Inf/NaN, no exceptions.
441/// Power operations use scalar fallback with logarithmic computation.
442#[inline(always)]
443pub fn float_masked_body_f64_simd<const LANES: usize>(
444    op: ArithmeticOperator,
445    lhs: &[f64],
446    rhs: &[f64],
447    mask: &Bitmask,
448    out: &mut [f64],
449    out_mask: &mut Bitmask,
450) {
451    type M = <f64 as SimdElement>::Mask;
452
453    let n = lhs.len();
454    let dense = all_true_mask_simd::<LANES>(mask);
455
456    if dense {
457        // hot
458        float_dense_body_f64_simd::<LANES>(op, lhs, rhs, out);
459        out_mask.fill(true);
460        return;
461    }
462
463    let mut i = 0;
464    while i + LANES <= n {
465        let a = Simd::<f64, LANES>::from_slice(&lhs[i..i + LANES]);
466        let b = Simd::<f64, LANES>::from_slice(&rhs[i..i + LANES]);
467        let m: Mask<M, LANES> = simd_mask::<M, LANES>(mask, i, n);
468
469        let res = match op {
470            ArithmeticOperator::Add => a + b,
471            ArithmeticOperator::Subtract => a - b,
472            ArithmeticOperator::Multiply => a * b,
473            ArithmeticOperator::Divide => a / b,
474            ArithmeticOperator::Remainder => a % b,
475            ArithmeticOperator::Power => (b * a.ln()).exp(),
476            ArithmeticOperator::FloorDiv => (a / b).floor(),
477        };
478
479        let selected = m.select(res, Simd::<f64, LANES>::splat(0.0));
480        selected.copy_to_slice(&mut out[i..i + LANES]);
481
482        write_simd_mask_bits(out_mask, i, m);
483        i += LANES;
484    }
485
486    // Tail often caused by `n % LANES =! 0`; uses scalar fallback
487    for j in i..n {
488        let valid = unsafe { mask.get_unchecked(j) };
489        if valid {
490            out[j] = match op {
491                ArithmeticOperator::Add => lhs[j] + rhs[j],
492                ArithmeticOperator::Subtract => lhs[j] - rhs[j],
493                ArithmeticOperator::Multiply => lhs[j] * rhs[j],
494                ArithmeticOperator::Divide => lhs[j] / rhs[j],
495                ArithmeticOperator::Remainder => lhs[j] % rhs[j],
496                ArithmeticOperator::Power => (rhs[j] * lhs[j].ln()).exp(),
497                ArithmeticOperator::FloorDiv => (lhs[j] / rhs[j]).floor(),
498            };
499            unsafe { out_mask.set_unchecked(j, true) };
500        } else {
501            out[j] = 0.0;
502            unsafe { out_mask.set_unchecked(j, false) };
503        }
504    }
505}
506
507/// SIMD f32 arithmetic kernel for dense arrays (no nulls).
508/// Vectorised operations with scalar fallback for power operations and array tails.
509/// Division by zero produces Inf/NaN following IEEE 754 semantics.
510#[inline(always)]
511pub fn float_dense_body_f32_simd<const LANES: usize>(
512    op: ArithmeticOperator,
513    lhs: &[f32],
514    rhs: &[f32],
515    out: &mut [f32],
516) {
517    let n = lhs.len();
518    let mut i = 0;
519    while i + LANES <= n {
520        let a = Simd::<f32, LANES>::from_slice(&lhs[i..i + LANES]);
521        let b = Simd::<f32, LANES>::from_slice(&rhs[i..i + LANES]);
522        let res = match op {
523            ArithmeticOperator::Add => a + b,
524            ArithmeticOperator::Subtract => a - b,
525            ArithmeticOperator::Multiply => a * b,
526            ArithmeticOperator::Divide => a / b,
527            ArithmeticOperator::Remainder => a % b,
528            ArithmeticOperator::Power => (b * a.ln()).exp(),
529            ArithmeticOperator::FloorDiv => (a / b).floor(),
530        };
531        res.copy_to_slice(&mut out[i..i + LANES]);
532        i += LANES;
533    }
534
535    // Tail often caused by `n % LANES =! 0`; uses scalar fallback
536    for j in i..n {
537        out[j] = match op {
538            ArithmeticOperator::Add => lhs[j] + rhs[j],
539            ArithmeticOperator::Subtract => lhs[j] - rhs[j],
540            ArithmeticOperator::Multiply => lhs[j] * rhs[j],
541            ArithmeticOperator::Divide => lhs[j] / rhs[j],
542            ArithmeticOperator::Remainder => lhs[j] % rhs[j],
543            ArithmeticOperator::Power => (rhs[j] * lhs[j].ln()).exp(),
544            ArithmeticOperator::FloorDiv => (lhs[j] / rhs[j]).floor(),
545        };
546    }
547}
548
549/// SIMD f64 arithmetic kernel for dense arrays (no nulls).
550/// Vectorised operations with scalar fallback for power operations and array tails.
551/// Division by zero produces Inf/NaN following IEEE 754 semantics.
552#[inline(always)]
553pub fn float_dense_body_f64_simd<const LANES: usize>(
554    op: ArithmeticOperator,
555    lhs: &[f64],
556    rhs: &[f64],
557    out: &mut [f64],
558) {
559    let n = lhs.len();
560    let mut i = 0;
561    while i + LANES <= n {
562        let a = Simd::<f64, LANES>::from_slice(&lhs[i..i + LANES]);
563        let b = Simd::<f64, LANES>::from_slice(&rhs[i..i + LANES]);
564        let res = match op {
565            ArithmeticOperator::Add => a + b,
566            ArithmeticOperator::Subtract => a - b,
567            ArithmeticOperator::Multiply => a * b,
568            ArithmeticOperator::Divide => a / b,
569            ArithmeticOperator::Remainder => a % b,
570            ArithmeticOperator::Power => (b * a.ln()).exp(),
571            ArithmeticOperator::FloorDiv => (a / b).floor(),
572        };
573        res.copy_to_slice(&mut out[i..i + LANES]);
574        i += LANES;
575    }
576
577    // Tail often caused by `n % LANES =! 0`; uses scalar fallback
578    for j in i..n {
579        out[j] = match op {
580            ArithmeticOperator::Add => lhs[j] + rhs[j],
581            ArithmeticOperator::Subtract => lhs[j] - rhs[j],
582            ArithmeticOperator::Multiply => lhs[j] * rhs[j],
583            ArithmeticOperator::Divide => lhs[j] / rhs[j],
584            ArithmeticOperator::Remainder => lhs[j] % rhs[j],
585            ArithmeticOperator::Power => (rhs[j] * lhs[j].ln()).exp(),
586            ArithmeticOperator::FloorDiv => (lhs[j] / rhs[j]).floor(),
587        };
588    }
589}
590
591/// SIMD f32 fused multiply-add kernel with null mask support.
592/// Hardware-accelerated `a.mul_add(b, c)` with proper null propagation.
593#[inline(always)]
594pub fn fma_masked_body_f32_simd<const LANES: usize>(
595    lhs: &[f32],
596    rhs: &[f32],
597    acc: &[f32],
598    mask: &Bitmask,
599    out: &mut [f32],
600    out_mask: &mut Bitmask,
601) {
602    use core::simd::{Mask, Simd};
603
604    let n = lhs.len();
605    let mut i = 0;
606    let dense = all_true_mask_simd::<LANES>(mask);
607
608    if dense {
609        fma_dense_body_f32_simd::<LANES>(lhs, rhs, acc, out);
610        out_mask.fill(true);
611        return;
612    }
613
614    while i + LANES <= n {
615        let a = Simd::<f32, LANES>::from_slice(&lhs[i..i + LANES]);
616        let b = Simd::<f32, LANES>::from_slice(&rhs[i..i + LANES]);
617        let c = Simd::<f32, LANES>::from_slice(&acc[i..i + LANES]);
618        let m: Mask<i32, LANES> = simd_mask::<i32, LANES>(mask, i, n);
619
620        let res = a.mul_add(b, c);
621
622        let selected = m.select(res, Simd::<f32, LANES>::splat(0.0));
623        selected.copy_to_slice(&mut out[i..i + LANES]);
624
625        write_simd_mask_bits(out_mask, i, m);
626        i += LANES;
627    }
628
629    // Scalar tail
630
631    for j in i..n {
632        let valid = unsafe { mask.get_unchecked(j) };
633        if valid {
634            out[j] = lhs[j].mul_add(rhs[j], acc[j]);
635            unsafe { out_mask.set_unchecked(j, true) };
636        } else {
637            out[j] = 0.0;
638            unsafe { out_mask.set_unchecked(j, false) };
639        }
640    }
641}
642
643/// SIMD f64 fused multiply-add kernel with null mask support.
644/// Hardware-accelerated `a.mul_add(b, c)` with proper null propagation.
645#[inline(always)]
646pub fn fma_masked_body_f64_simd<const LANES: usize>(
647    lhs: &[f64],
648    rhs: &[f64],
649    acc: &[f64],
650    mask: &Bitmask,
651    out: &mut [f64],
652    out_mask: &mut Bitmask,
653) {
654    use core::simd::{Mask, Simd};
655
656    let n = lhs.len();
657    let mut i = 0;
658    let dense = all_true_mask_simd::<LANES>(mask);
659
660    if dense {
661        // Hot
662        fma_dense_body_f64_simd::<LANES>(lhs, rhs, acc, out);
663        out_mask.fill(true);
664        return;
665    }
666
667    while i + LANES <= n {
668        let a = Simd::<f64, LANES>::from_slice(&lhs[i..i + LANES]);
669        let b = Simd::<f64, LANES>::from_slice(&rhs[i..i + LANES]);
670        let c = Simd::<f64, LANES>::from_slice(&acc[i..i + LANES]);
671        let m: Mask<i64, LANES> = simd_mask::<i64, LANES>(mask, i, n);
672
673        let res = a.mul_add(b, c);
674
675        let selected = m.select(res, Simd::<f64, LANES>::splat(0.0));
676        selected.copy_to_slice(&mut out[i..i + LANES]);
677
678        write_simd_mask_bits(out_mask, i, m);
679        i += LANES;
680    }
681
682    // Scalar tail
683
684    for j in i..n {
685        let valid = unsafe { mask.get_unchecked(j) };
686        if valid {
687            out[j] = lhs[j].mul_add(rhs[j], acc[j]);
688            unsafe { out_mask.set_unchecked(j, true) };
689        } else {
690            out[j] = 0.0;
691            unsafe { out_mask.set_unchecked(j, false) };
692        }
693    }
694}
695
696/// SIMD f32 fused multiply-add kernel for dense arrays (no nulls).
697/// Hardware-accelerated `a.mul_add(b, c)` with vectorised and scalar tail processing.
698#[inline(always)]
699pub fn fma_dense_body_f32_simd<const LANES: usize>(
700    lhs: &[f32],
701    rhs: &[f32],
702    acc: &[f32],
703    out: &mut [f32],
704) {
705    use core::simd::Simd;
706
707    let n = lhs.len();
708    let mut i = 0;
709
710    while i + LANES <= n {
711        let a = Simd::<f32, LANES>::from_slice(&lhs[i..i + LANES]);
712        let b = Simd::<f32, LANES>::from_slice(&rhs[i..i + LANES]);
713        let c = Simd::<f32, LANES>::from_slice(&acc[i..i + LANES]);
714        let res = a.mul_add(b, c);
715        res.copy_to_slice(&mut out[i..i + LANES]);
716        i += LANES;
717    }
718
719    for j in i..n {
720        out[j] = lhs[j].mul_add(rhs[j], acc[j]);
721    }
722}
723
724/// SIMD f64 fused multiply-add kernel for dense arrays (no nulls).
725/// Hardware-accelerated `a.mul_add(b, c)` with vectorised and scalar tail processing.
726#[inline(always)]
727pub fn fma_dense_body_f64_simd<const LANES: usize>(
728    lhs: &[f64],
729    rhs: &[f64],
730    acc: &[f64],
731    out: &mut [f64],
732) {
733    use core::simd::Simd;
734
735    let n = lhs.len();
736    let mut i = 0;
737
738    while i + LANES <= n {
739        let a = Simd::<f64, LANES>::from_slice(&lhs[i..i + LANES]);
740        let b = Simd::<f64, LANES>::from_slice(&rhs[i..i + LANES]);
741        let c = Simd::<f64, LANES>::from_slice(&acc[i..i + LANES]);
742        let res = a.mul_add(b, c);
743        res.copy_to_slice(&mut out[i..i + LANES]);
744        i += LANES;
745    }
746
747    // Tail uses scalar fallback
748    for j in i..n {
749        out[j] = lhs[j].mul_add(rhs[j], acc[j]);
750    }
751}
simd_kernels/kernels/arithmetic/simd.rs

simd_kernels/kernels/arithmetic/
simd.rs
