Skip to main content

ruda_kernel/dsl/frontend/container/vector/
ops.rs

1use core::{marker::PhantomData, ops::Not};
2use ruda_core::ir::{Bitwise, ConstantValue, ElemType, Instruction, Type, UIntKind, UnaryOperator};
3use ruda_kernel_macros::{ruda, intrinsic};
4use num_traits::{NumCast, One, ToPrimitive, Zero};
5
6use crate::dsl::{
7    prelude::{
8        ArcTan2, InverseSqrt, IsInf, IsNan, Powf, Powi, SaturatingAdd, SaturatingSub, Trunc,
9    },
10};
11use crate::dsl::{prelude::*, unexpanded};
12
13use super::Vector;
14type VectorExpand<E, N> = NativeExpand<Vector<E, N>>;
15
16impl<P, N: Size> core::ops::Add<Self> for Vector<P, N>
17where
18    P: Scalar,
19    P: core::ops::Add<P, Output = P>,
20{
21    type Output = Self;
22
23    fn add(self, rhs: Self) -> Self::Output {
24        Self::new(self.val + rhs.val)
25    }
26}
27
28impl<P, N: Size> core::ops::Sub<Self> for Vector<P, N>
29where
30    P: Scalar,
31    P: core::ops::Sub<P, Output = P>,
32{
33    type Output = Self;
34
35    fn sub(self, rhs: Self) -> Self::Output {
36        Self::new(self.val - rhs.val)
37    }
38}
39
40impl<P, N: Size> core::ops::Mul<Self> for Vector<P, N>
41where
42    P: Scalar,
43    P: core::ops::Mul<P, Output = P>,
44{
45    type Output = Self;
46
47    fn mul(self, rhs: Self) -> Self::Output {
48        Self::new(self.val * rhs.val)
49    }
50}
51
52impl<P, N: Size> core::ops::Div<Self> for Vector<P, N>
53where
54    P: Scalar,
55    P: core::ops::Div<P, Output = P>,
56{
57    type Output = Self;
58
59    fn div(self, rhs: Self) -> Self::Output {
60        Self::new(self.val / rhs.val)
61    }
62}
63
64impl<P, N: Size> core::ops::AddAssign<Self> for Vector<P, N>
65where
66    P: Scalar,
67    P: core::ops::AddAssign,
68{
69    fn add_assign(&mut self, rhs: Self) {
70        self.val += rhs.val;
71    }
72}
73
74impl<P, N: Size> core::ops::SubAssign<Self> for Vector<P, N>
75where
76    P: Scalar,
77    P: core::ops::SubAssign,
78{
79    fn sub_assign(&mut self, rhs: Self) {
80        self.val -= rhs.val;
81    }
82}
83
84impl<P, N: Size> core::ops::DivAssign<Self> for Vector<P, N>
85where
86    P: Scalar,
87    P: core::ops::DivAssign,
88{
89    fn div_assign(&mut self, rhs: Self) {
90        self.val /= rhs.val;
91    }
92}
93
94impl<P, N: Size> core::ops::MulAssign<Self> for Vector<P, N>
95where
96    P: Scalar,
97    P: core::ops::MulAssign,
98{
99    fn mul_assign(&mut self, rhs: Self) {
100        self.val *= rhs.val;
101    }
102}
103
104impl<P, N: Size> core::cmp::PartialEq for Vector<P, N>
105where
106    P: Scalar,
107    P: core::cmp::PartialEq,
108{
109    fn eq(&self, other: &Self) -> bool {
110        self.val.eq(&other.val)
111    }
112}
113
114impl<P, N: Size> core::cmp::PartialOrd for Vector<P, N>
115where
116    P: Scalar,
117    P: core::cmp::PartialOrd,
118{
119    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
120        self.val.partial_cmp(&other.val)
121    }
122}
123
124impl<P, N: Size> core::ops::BitAnd<Self> for Vector<P, N>
125where
126    P: Scalar,
127    P: core::ops::BitAnd<P, Output = P>,
128{
129    type Output = Self;
130
131    fn bitand(self, rhs: Self) -> Self::Output {
132        Self::new(self.val & rhs.val)
133    }
134}
135
136impl<P, N: Size> core::ops::BitOr<Self> for Vector<P, N>
137where
138    P: Scalar,
139    P: core::ops::BitOr<P, Output = P>,
140{
141    type Output = Self;
142
143    fn bitor(self, rhs: Self) -> Self::Output {
144        Self::new(self.val | rhs.val)
145    }
146}
147
148impl<P, N: Size> core::ops::BitXor<Self> for Vector<P, N>
149where
150    P: Scalar,
151    P: core::ops::BitXor<P, Output = P>,
152{
153    type Output = Self;
154
155    fn bitxor(self, rhs: Self) -> Self::Output {
156        Self::new(self.val ^ rhs.val)
157    }
158}
159
160impl<P, N: Size> core::ops::Shl<Self> for Vector<P, N>
161where
162    P: Scalar,
163    P: core::ops::Shl<P, Output = P>,
164{
165    type Output = Self;
166
167    fn shl(self, rhs: Self) -> Self::Output {
168        Self::new(self.val << rhs.val)
169    }
170}
171
172impl<P, N: Size> core::ops::Shr<Self> for Vector<P, N>
173where
174    P: Scalar,
175    P: core::ops::Shr<P, Output = P>,
176{
177    type Output = Self;
178
179    fn shr(self, rhs: Self) -> Self::Output {
180        Self::new(self.val >> rhs.val)
181    }
182}
183
184impl<P, N: Size> core::ops::BitAndAssign<Self> for Vector<P, N>
185where
186    P: Scalar,
187    P: core::ops::BitAndAssign,
188{
189    fn bitand_assign(&mut self, rhs: Self) {
190        self.val &= rhs.val;
191    }
192}
193
194impl<P, N: Size> core::ops::BitOrAssign<Self> for Vector<P, N>
195where
196    P: Scalar,
197    P: core::ops::BitOrAssign,
198{
199    fn bitor_assign(&mut self, rhs: Self) {
200        self.val |= rhs.val;
201    }
202}
203
204impl<P, N: Size> core::ops::BitXorAssign<Self> for Vector<P, N>
205where
206    P: Scalar,
207    P: core::ops::BitXorAssign,
208{
209    fn bitxor_assign(&mut self, rhs: Self) {
210        self.val ^= rhs.val;
211    }
212}
213
214impl<P, N: Size> core::ops::ShlAssign<Self> for Vector<P, N>
215where
216    P: Scalar,
217    P: core::ops::ShlAssign,
218{
219    fn shl_assign(&mut self, rhs: Self) {
220        self.val <<= rhs.val;
221    }
222}
223
224impl<P, N: Size> core::ops::ShrAssign<Self> for Vector<P, N>
225where
226    P: Scalar,
227    P: core::ops::ShrAssign,
228{
229    fn shr_assign(&mut self, rhs: Self) {
230        self.val >>= rhs.val;
231    }
232}
233
234impl<P: Scalar + Abs, N: Size> Abs for Vector<P, N> {}
235impl<P: Scalar + Log, N: Size> Log for Vector<P, N> {}
236impl<P: Scalar + Log1p, N: Size> Log1p for Vector<P, N> {}
237impl<P: Scalar + Erf, N: Size> Erf for Vector<P, N> {}
238impl<P: Scalar + Exp, N: Size> Exp for Vector<P, N> {}
239impl<P: Scalar + Powf, N: Size> Powf for Vector<P, N> {}
240impl<P: Scalar + Powi<I>, I: Scalar, N: Size> Powi<Vector<I, N>> for Vector<P, N> {}
241impl<P: Scalar + Sqrt, N: Size> Sqrt for Vector<P, N> {}
242impl<P: Scalar + InverseSqrt, N: Size> InverseSqrt for Vector<P, N> {}
243impl<P: Scalar + Cos, N: Size> Cos for Vector<P, N> {}
244impl<P: Scalar + Sin, N: Size> Sin for Vector<P, N> {}
245impl<P: Scalar + Tan, N: Size> Tan for Vector<P, N> {}
246impl<P: Scalar + Tanh, N: Size> Tanh for Vector<P, N> {}
247impl<P: Scalar + Sinh, N: Size> Sinh for Vector<P, N> {}
248impl<P: Scalar + Cosh, N: Size> Cosh for Vector<P, N> {}
249impl<P: Scalar + ArcSin, N: Size> ArcSin for Vector<P, N> {}
250impl<P: Scalar + ArcCos, N: Size> ArcCos for Vector<P, N> {}
251impl<P: Scalar + ArcTan, N: Size> ArcTan for Vector<P, N> {}
252impl<P: Scalar + ArcSinh, N: Size> ArcSinh for Vector<P, N> {}
253impl<P: Scalar + ArcCosh, N: Size> ArcCosh for Vector<P, N> {}
254impl<P: Scalar + ArcTanh, N: Size> ArcTanh for Vector<P, N> {}
255impl<P: Scalar + ArcTan2, N: Size> ArcTan2 for Vector<P, N> {}
256impl<P: Scalar + Recip, N: Size> Recip for Vector<P, N> {}
257impl<P: Scalar + Remainder, N: Size> Remainder for Vector<P, N> {}
258impl<P: Scalar + Round, N: Size> Round for Vector<P, N> {}
259impl<P: Scalar + Floor, N: Size> Floor for Vector<P, N> {}
260impl<P: Scalar + Ceil, N: Size> Ceil for Vector<P, N> {}
261impl<P: Scalar + Trunc, N: Size> Trunc for Vector<P, N> {}
262impl<P: Scalar + ReverseBits, N: Size> ReverseBits for Vector<P, N> {}
263impl<P: Scalar + RudaNot, N: Size> RudaNot for Vector<P, N> {}
264impl<P: Scalar + SaturatingAdd, N: Size> SaturatingAdd for Vector<P, N> {}
265impl<P: Scalar + SaturatingSub, N: Size> SaturatingSub for Vector<P, N> {}
266impl<P: Scalar + IsNan, N: Size> IsNan for Vector<P, N> {}
267impl<P: Scalar + IsInf, N: Size> IsInf for Vector<P, N> {}
268impl<P: Scalar + Normalize, N: Size> Normalize for Vector<P, N> {}
269impl<P: Scalar + Magnitude, N: Size> Magnitude for Vector<P, N> {}
270impl<P: Scalar + VectorSum, N: Size> VectorSum for Vector<P, N> {}
271impl<P: Scalar + Degrees, N: Size> Degrees for Vector<P, N> {}
272impl<P: Scalar + Radians, N: Size> Radians for Vector<P, N> {}
273
274impl<P: Scalar + Ord, N: Size> Ord for Vector<P, N> {
275    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
276        self.val.cmp(&other.val)
277    }
278}
279
280#[ruda]
281impl<P: CountOnes + Scalar, N: Size> Vector<P, N> {
282    pub fn count_ones(self) -> Vector<u32, N> {
283        intrinsic!(|scope| {
284            let out_item = Type::scalar(ElemType::UInt(UIntKind::U32))
285                .with_vector_size(self.expand.ty.vector_size());
286            let out = scope.create_local(out_item);
287            scope.register(Instruction::new(
288                Bitwise::CountOnes(UnaryOperator {
289                    input: *self.expand,
290                }),
291                *out,
292            ));
293            out.into()
294        })
295    }
296}
297
298impl<P: LeadingZeros + Scalar, N: Size> LeadingZeros for Vector<P, N> {}
299impl<P: FindFirstSet + Scalar, N: Size> FindFirstSet for Vector<P, N> {}
300impl<P: TrailingZeros + Scalar, N: Size> TrailingZeros for Vector<P, N> {}
301
302impl<P: Scalar + NumCast, N: Size> NumCast for Vector<P, N> {
303    fn from<T: num_traits::ToPrimitive>(n: T) -> Option<Self> {
304        let val: P = NumCast::from(n)?;
305        Some(Self {
306            val,
307            _size: PhantomData,
308        })
309    }
310}
311impl<P: Scalar + NumCast, N: Size> ToPrimitive for Vector<P, N> {
312    fn to_i64(&self) -> Option<i64> {
313        self.val.to_i64()
314    }
315
316    fn to_u64(&self) -> Option<u64> {
317        self.val.to_u64()
318    }
319}
320
321impl<P: Not<Output = P> + Scalar, N: Size> Not for Vector<P, N> {
322    type Output = Self;
323
324    fn not(self) -> Self::Output {
325        Vector::new(self.val.not())
326    }
327}
328
329#[allow(clippy::from_over_into)]
330impl<P: Scalar + Into<NativeExpand<P>>, N: Size> Into<NativeExpand<Self>> for Vector<P, N> {
331    fn into(self) -> NativeExpand<Self> {
332        let elem: NativeExpand<P> = self.val.into();
333        elem.expand.into()
334    }
335}
336
337impl<T: Scalar + Default, N: Size> Default for Vector<T, N> {
338    fn default() -> Self {
339        Self::new(T::default())
340    }
341}
342
343impl<T: Scalar + IntoRuntime, N: Size> IntoRuntime for Vector<T, N> {
344    fn __expand_runtime_method(self, scope: &mut Scope) -> Self::ExpandType {
345        let val = self.val.__expand_runtime_method(scope);
346        Self::__expand_new(scope, val)
347    }
348}
349
350impl<T: Scalar + Into<ConstantValue>, N: Size> From<Vector<T, N>> for ConstantValue {
351    fn from(value: Vector<T, N>) -> Self {
352        value.val.into()
353    }
354}
355
356impl<T: Scalar + Zero, N: Size> Zero for Vector<T, N> {
357    fn zero() -> Self {
358        Self::new(T::zero())
359    }
360
361    fn is_zero(&self) -> bool {
362        self.val.is_zero()
363    }
364}
365
366impl<T: Scalar + One, N: Size> One for Vector<T, N> {
367    fn one() -> Self {
368        Self::new(T::one())
369    }
370}
371
372macro_rules! operation_literal {
373    ($lit:ty) => {
374        impl<P, N: Size> core::ops::Add<$lit> for Vector<P, N>
375        where
376            P: Scalar,
377        {
378            type Output = Self;
379
380            fn add(self, _rhs: $lit) -> Self::Output {
381                unexpanded!();
382            }
383        }
384
385        impl<P, N: Size> core::ops::Sub<$lit> for Vector<P, N>
386        where
387            P: Scalar,
388        {
389            type Output = Self;
390
391            fn sub(self, _rhs: $lit) -> Self::Output {
392                unexpanded!();
393            }
394        }
395
396        impl<P, N: Size> core::ops::Mul<$lit> for Vector<P, N>
397        where
398            P: Scalar,
399        {
400            type Output = Self;
401
402            fn mul(self, _rhs: $lit) -> Self::Output {
403                unexpanded!();
404            }
405        }
406
407        impl<P, N: Size> core::ops::Div<$lit> for Vector<P, N>
408        where
409            P: Scalar,
410        {
411            type Output = Self;
412
413            fn div(self, _rhs: $lit) -> Self::Output {
414                unexpanded!();
415            }
416        }
417    };
418}
419
420operation_literal!(f32);
421operation_literal!(f64);
422operation_literal!(usize);
423operation_literal!(i32);
424operation_literal!(i64);