Skip to main content

simd_kernels/kernels/arithmetic/
std.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//! # **Standard Arithmetic Kernels Module** - *Scalar Fallback / Non-SIMD Implementations*
16//!
17//! Portable scalar implementations of arithmetic operations for compatibility and unaligned data.
18//!
19//! Prefer dispatch.rs for easily handling the general case, otherwise you can use these inner functions
20//! directly (e.g., "dense_std") vs. "maybe masked, maybe std".
21//!
22//! ## Overview
23//! - **Scalar loops**: Standard element-wise operations without vectorisation
24//! - **Fallback role**: Used when SIMD alignment requirements aren't met or SIMD is disabled
25//! - **Full compatibility**: Works on any architecture regardless of SIMD support
26//! - **Null-aware**: Supports Arrow-compatible null mask propagation
27//!
28//! ## Design Notes
29//! - Intentionally avoids parallelisation to allow higher-level chunking strategies
30//! - Wrapping arithmetic for integers to prevent overflow panics
31//! - Division by zero handling: panics for integers, produces Inf/NaN for floats
32
33use crate::Bitmask;
34use crate::enums::operators::ArithmeticOperator;
35use num_traits::{Float, PrimInt, ToPrimitive, WrappingAdd, WrappingMul, WrappingSub};
36
37/// Scalar integer arithmetic kernel for dense arrays (no nulls).
38/// Performs element-wise operations using wrapping arithmetic to prevent overflow panics.
39/// Panics on division/remainder by zero.
40#[inline(always)]
41pub fn int_dense_body_std<T: PrimInt + ToPrimitive + WrappingAdd + WrappingSub + WrappingMul>(
42    op: ArithmeticOperator,
43    lhs: &[T],
44    rhs: &[T],
45    out: &mut [T],
46) {
47    let n = lhs.len();
48    for i in 0..n {
49        out[i] = match op {
50            ArithmeticOperator::Add => lhs[i].wrapping_add(&rhs[i]),
51            ArithmeticOperator::Subtract => lhs[i].wrapping_sub(&rhs[i]),
52            ArithmeticOperator::Multiply => lhs[i].wrapping_mul(&rhs[i]),
53            ArithmeticOperator::Divide => {
54                if rhs[i] == T::zero() {
55                    panic!("Division by zero")
56                } else {
57                    lhs[i] / rhs[i]
58                }
59            }
60            ArithmeticOperator::Remainder => {
61                if rhs[i] == T::zero() {
62                    panic!("Remainder by zero")
63                } else {
64                    lhs[i] % rhs[i]
65                }
66            }
67            ArithmeticOperator::Power => lhs[i].pow(rhs[i].to_u32().unwrap_or(0)),
68            ArithmeticOperator::FloorDiv => {
69                if rhs[i] == T::zero() {
70                    panic!("Floor division by zero")
71                } else {
72                    let d = lhs[i] / rhs[i];
73                    let r = lhs[i] % rhs[i];
74                    // If remainder is non-zero and signs differ, floor toward -inf
75                    if r != T::zero() && (lhs[i] ^ rhs[i]) < T::zero() { d - T::one() } else { d }
76                }
77            }
78        };
79    }
80}
81
82/// Scalar integer arithmetic kernel with null mask support.
83/// Handles division by zero gracefully by marking results as null instead of panicking.
84/// Invalid inputs (mask=false) and zero division produce null outputs.
85#[inline(always)]
86pub fn int_masked_body_std<T: PrimInt + ToPrimitive + WrappingAdd + WrappingSub + WrappingMul>(
87    op: ArithmeticOperator,
88    lhs: &[T],
89    rhs: &[T],
90    mask: &Bitmask,
91    out: &mut [T],
92    out_mask: &mut Bitmask,
93) {
94    let n = lhs.len();
95    for i in 0..n {
96        let valid = unsafe { mask.get_unchecked(i) };
97        if valid {
98            let (result, final_valid) = match op {
99                ArithmeticOperator::Add => (lhs[i].wrapping_add(&rhs[i]), true),
100                ArithmeticOperator::Subtract => (lhs[i].wrapping_sub(&rhs[i]), true),
101                ArithmeticOperator::Multiply => (lhs[i].wrapping_mul(&rhs[i]), true),
102                ArithmeticOperator::Divide => {
103                    if rhs[i] == T::zero() {
104                        (T::zero(), false) // division by zero -> invalid
105                    } else {
106                        (lhs[i] / rhs[i], true)
107                    }
108                }
109                ArithmeticOperator::Remainder => {
110                    if rhs[i] == T::zero() {
111                        (T::zero(), false) // remainder by zero -> invalid
112                    } else {
113                        (lhs[i] % rhs[i], true)
114                    }
115                }
116                ArithmeticOperator::Power => (lhs[i].pow(rhs[i].to_u32().unwrap_or(0)), true),
117                ArithmeticOperator::FloorDiv => {
118                    if rhs[i] == T::zero() {
119                        (T::zero(), false)
120                    } else {
121                        let d = lhs[i] / rhs[i];
122                        let r = lhs[i] % rhs[i];
123                        if r != T::zero() && (lhs[i] ^ rhs[i]) < T::zero() { (d - T::one(), true) } else { (d, true) }
124                    }
125                }
126            };
127            out[i] = result;
128            unsafe {
129                out_mask.set_unchecked(i, final_valid);
130            }
131        } else {
132            out[i] = T::zero();
133            unsafe {
134                out_mask.set_unchecked(i, false);
135            }
136        }
137    }
138}
139
140/// Scalar floating-point arithmetic kernel for dense arrays (no nulls).
141/// Division by zero produces Inf/NaN rather than panicking.
142/// Power operations use logarithmic exponentiation: `exp(b * ln(a))`.
143#[inline(always)]
144pub fn float_dense_body_std<T: Float>(op: ArithmeticOperator, lhs: &[T], rhs: &[T], out: &mut [T]) {
145    let n = lhs.len();
146    for i in 0..n {
147        out[i] = match op {
148            ArithmeticOperator::Add => lhs[i] + rhs[i],
149            ArithmeticOperator::Subtract => lhs[i] - rhs[i],
150            ArithmeticOperator::Multiply => lhs[i] * rhs[i],
151            ArithmeticOperator::Divide => lhs[i] / rhs[i],
152            ArithmeticOperator::Remainder => lhs[i] % rhs[i],
153            ArithmeticOperator::Power => (rhs[i] * lhs[i].ln()).exp(),
154            ArithmeticOperator::FloorDiv => (lhs[i] / rhs[i]).floor(),
155        };
156    }
157}
158
159/// Scalar floating-point arithmetic kernel with null mask support.
160/// Preserves IEEE 754 semantics: division by zero produces Inf/NaN, no panicking.
161/// Invalid inputs (mask=false) produce null outputs with zero values.
162#[inline(always)]
163pub fn float_masked_body_std<T: Float>(
164    op: ArithmeticOperator,
165    lhs: &[T],
166    rhs: &[T],
167    mask: &Bitmask,
168    out: &mut [T],
169    out_mask: &mut Bitmask,
170) {
171    let n = lhs.len();
172    for i in 0..n {
173        let valid = unsafe { mask.get_unchecked(i) };
174        if valid {
175            out[i] = match op {
176                ArithmeticOperator::Add => lhs[i] + rhs[i],
177                ArithmeticOperator::Subtract => lhs[i] - rhs[i],
178                ArithmeticOperator::Multiply => lhs[i] * rhs[i],
179                ArithmeticOperator::Divide => lhs[i] / rhs[i],
180                ArithmeticOperator::Remainder => lhs[i] % rhs[i],
181                ArithmeticOperator::Power => (rhs[i] * lhs[i].ln()).exp(),
182                ArithmeticOperator::FloorDiv => (lhs[i] / rhs[i]).floor(),
183            };
184            unsafe {
185                out_mask.set_unchecked(i, true);
186            }
187        } else {
188            out[i] = T::zero();
189            unsafe {
190                out_mask.set_unchecked(i, false);
191            }
192        }
193    }
194}
195
196/// Fused multiply add (a * b + acc) with null mask
197#[inline(always)]
198pub fn fma_masked_body_std<T: Float>(
199    lhs: &[T],
200    rhs: &[T],
201    acc: &[T],
202    mask: &Bitmask,
203    out: &mut [T],
204    out_mask: &mut Bitmask,
205) {
206    let n = lhs.len();
207    for i in 0..n {
208        let valid = unsafe { mask.get_unchecked(i) };
209        if valid {
210            out[i] = lhs[i].mul_add(rhs[i], acc[i]);
211            unsafe {
212                out_mask.set_unchecked(i, true);
213            }
214        } else {
215            out[i] = T::zero();
216            unsafe {
217                out_mask.set_unchecked(i, false);
218            }
219        }
220    }
221}
222
223/// Dense fused multiply add (a * b + acc)
224#[inline(always)]
225pub fn fma_dense_body_std<T: Float>(lhs: &[T], rhs: &[T], acc: &[T], out: &mut [T]) {
226    let n = lhs.len();
227    for i in 0..n {
228        out[i] = lhs[i].mul_add(rhs[i], acc[i]);
229    }
230}