simd_kernels/kernels/arithmetic/
std.rs1use crate::Bitmask;
34use crate::enums::operators::ArithmeticOperator;
35use num_traits::{Float, PrimInt, ToPrimitive, WrappingAdd, WrappingMul, WrappingSub};
36
37#[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 r != T::zero() && (lhs[i] ^ rhs[i]) < T::zero() { d - T::one() } else { d }
76 }
77 }
78 };
79 }
80}
81
82#[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) } else {
106 (lhs[i] / rhs[i], true)
107 }
108 }
109 ArithmeticOperator::Remainder => {
110 if rhs[i] == T::zero() {
111 (T::zero(), false) } 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#[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#[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#[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#[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}