cubecl_core/frontend/container/vector/
ops.rs

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

cubecl_core/frontend/container/vector/
ops.rs
