cubecl_core/frontend/container/vector/
base.rs

1use core::{marker::PhantomData, ops::Neg};
2
3use crate::frontend::{CubePrimitive, CubeType, NativeAssign, NativeExpand};
4use crate::ir::{BinaryOperands, Instruction, Scope, Type};
5use crate::{self as cubecl, prelude::*};
6use cubecl_ir::{Comparison, ConstantValue, Value};
7use cubecl_macros::{cube, intrinsic};
8
9/// A contiguous list of elements that supports auto-vectorized operations.
10#[derive(Debug)]
11pub struct Vector<P: Scalar, N: Size> {
12    // Comptime vectors only support 1 element.
13    pub(crate) val: P,
14    pub(crate) _size: PhantomData<N>,
15}
16
17type VectorExpand<P, N> = NativeExpand<Vector<P, N>>;
18
19impl<P: Scalar, N: Size> Clone for Vector<P, N> {
20    fn clone(&self) -> Self {
21        *self
22    }
23}
24impl<P: Scalar, N: Size> Eq for Vector<P, N> {}
25impl<P: Scalar, N: Size> Copy for Vector<P, N> {}
26impl<P: Scalar + Neg<Output = P>, N: Size> Neg for Vector<P, N> {
27    type Output = Self;
28
29    fn neg(self) -> Self::Output {
30        Self {
31            val: -self.val,
32            _size: PhantomData,
33        }
34    }
35}
36
37/// Module that contains the implementation details of the new function.
38mod new {
39    use cubecl_ir::VectorSize;
40    use cubecl_macros::comptime_type;
41
42    use crate::prelude::Cast;
43
44    use super::*;
45
46    impl<P: Scalar, N: Size> Vector<P, N> {
47        /// Create a new vector of size 1 using the given value.
48        #[allow(unused_variables)]
49        pub fn new(val: P) -> Self {
50            Self {
51                val,
52                _size: PhantomData,
53            }
54        }
55
56        pub fn __expand_new(scope: &Scope, val: NativeExpand<P>) -> VectorExpand<P, N> {
57            Vector::<P, N>::__expand_cast_from(scope, val)
58        }
59    }
60
61    impl<P: Scalar, N: Size> Vector<P, N> {
62        /// Get the length of the current vector.
63        pub fn vector_size(&self) -> comptime_type!(VectorSize) {
64            N::value()
65        }
66    }
67}
68
69mod components {
70    use cubecl_ir::{Operator, VectorInsertOperands};
71
72    use super::*;
73
74    #[cube]
75    impl<P: Scalar, N: Size> Vector<P, N> {
76        #[allow(unused)]
77        pub fn extract(self, index: usize) -> P {
78            intrinsic!(|scope| {
79                if self.expand.vector_size() > 1 {
80                    let out = scope.create_value(P::__expand_as_type(scope));
81                    scope.register(Instruction::new(
82                        Operator::ExtractComponent(BinaryOperands {
83                            lhs: self.expand,
84                            rhs: index.expand,
85                        }),
86                        out,
87                    ));
88                    out.into()
89                } else {
90                    read_value(scope, self.expand).into()
91                }
92            })
93        }
94
95        #[allow(unused)]
96        pub fn insert(&mut self, index: usize, value: P) {
97            intrinsic!(|scope| {
98                if self.expand.vector_size() > 1 {
99                    let new_value = scope.create_value(self.expand.ty.unwrap_ptr());
100                    scope.register(Instruction::new(
101                        Operator::InsertComponent(VectorInsertOperands {
102                            vector: self.expand,
103                            index: index.expand,
104                            value: value.expand,
105                        }),
106                        new_value,
107                    ));
108                    assign::expand_element(scope, new_value, self.expand);
109                } else {
110                    assign::expand_element(scope, value.expand, self.expand);
111                }
112            })
113        }
114    }
115}
116
117mod numeric {
118    use super::*;
119
120    #[cube]
121    impl<P: Numeric, N: Size> Vector<P, N> {
122        pub fn min_value() -> Self {
123            Self::new(P::min_value())
124        }
125        pub fn max_value() -> Self {
126            Self::new(P::max_value())
127        }
128
129        /// Create a new constant numeric.
130        ///
131        /// Note: since this must work for both integer and float
132        /// only the less expressive of both can be created (int)
133        /// If a number with decimals is needed, use `Float::new`.
134        ///
135        /// This method panics when unexpanded. For creating an element
136        /// with a val, use the new method of the sub type.
137        pub fn from_int(val: i64) -> Self {
138            Self::new(P::from_int(val))
139        }
140    }
141}
142
143/// Module that contains the implementation details of the fill function.
144mod fill {
145    use super::*;
146
147    #[cube]
148    impl<P: Scalar, N: Size> Vector<P, N> {
149        /// Fill the vector with the given value.
150        ///
151        /// If you want to fill the vector with different values, consider using the index API
152        /// instead.
153        ///
154        /// ```rust, ignore
155        /// let mut vector = Vector::<u32>::empty(2);
156        /// vector[0] = 1;
157        /// vector[1] = 2;
158        /// ```
159        pub fn fill(self, value: P) -> Self {
160            intrinsic!(|scope| { Vector::<P, N>::__expand_cast_from(scope, value) })
161        }
162    }
163}
164
165/// Module that contains the implementation details of the empty function.
166mod empty {
167    use bytemuck::Zeroable;
168
169    use super::*;
170
171    #[cube]
172    impl<P: Scalar, N: Size> Vector<P, N> {
173        pub fn empty() -> Self {
174            intrinsic!(|scope| {
175                let value = Self::__expand_default(scope);
176                value.into_mut(scope)
177            })
178        }
179    }
180
181    #[cube]
182    impl<P: Scalar + Zeroable, N: Size> Vector<P, N> {
183        pub fn zeroed() -> Self {
184            intrinsic!(|scope| {
185                let zeroed = P::zeroed().__expand_runtime_method(scope);
186                Self::__expand_cast_from(scope, zeroed)
187            })
188        }
189    }
190}
191
192/// Module that contains the implementation details of the size function.
193mod size {
194    use cubecl_ir::VectorSize;
195
196    use super::*;
197
198    impl<P: Scalar, N: Size> Vector<P, N> {
199        /// Get the number of individual elements a vector contains.
200        ///
201        /// The size is available at comptime and may be used in combination with the comptime
202        /// macro.
203        ///
204        /// ```rust, ignore
205        /// // The if statement is going to be executed at comptime.
206        /// if comptime!(vector.size() == 1) {
207        /// }
208        /// ```
209        pub fn size(&self) -> VectorSize {
210            N::value()
211        }
212
213        /// Expand function of [size](Self::size).
214        pub fn __expand_size(scope: &Scope, element: NativeExpand<Vector<P, N>>) -> VectorSize {
215            element.__expand_vector_size_method(scope)
216        }
217    }
218
219    impl<P: Scalar, N: Size> NativeExpand<Vector<P, N>> {
220        /// Comptime version of [size](Vector::size).
221        pub fn size(&self) -> VectorSize {
222            self.expand.ty.vector_size()
223        }
224
225        /// Expand method of [size](Vector::size).
226        pub fn __expand_size_method(&self, _scope: &Scope) -> VectorSize {
227            self.size()
228        }
229    }
230}
231
232// Implement a comparison operator define in
233macro_rules! impl_vector_comparison {
234    ($name:ident, $operator:ident, $comment:literal) => {
235        ::paste::paste! {
236            /// Module that contains the implementation details of the $name function.
237            mod $name {
238
239                use super::*;
240
241                #[cube]
242                impl<P: Scalar, N: Size> Vector<P, N> {
243                    #[doc = concat!(
244                        "Return a new vector with the element-wise comparison of the first vector being ",
245                        $comment,
246                        " the second vector."
247                    )]
248                    pub fn $name(&self, other: &Self) -> Vector<bool, N> {
249                        intrinsic!(|scope| {
250                            let this = self.__expand_deref_method(scope);
251                            let other = other.__expand_deref_method(scope);
252                            let size = this.expand.ty.vector_size();
253                            let lhs = this.expand;
254                            let rhs = other.expand.into();
255
256                            let output = scope.create_value(Vector::<bool, N>::__expand_as_type(scope));
257
258                            scope.register(Instruction::new(
259                                Comparison::$operator(BinaryOperands { lhs, rhs }),
260                                output.clone().into(),
261                            ));
262
263                            output.into()
264                        })
265                    }
266                }
267            }
268        }
269
270    };
271}
272
273impl_vector_comparison!(equal, Equal, "equal to");
274impl_vector_comparison!(not_equal, NotEqual, "not equal to");
275impl_vector_comparison!(less_than, Lower, "less than");
276impl_vector_comparison!(greater_than, Greater, "greater than");
277impl_vector_comparison!(less_equal, LowerEqual, "less than or equal to");
278impl_vector_comparison!(greater_equal, GreaterEqual, "greater than or equal to");
279
280mod bool_and {
281    use cubecl_ir::Operator;
282
283    use crate::prelude::binary_expand;
284
285    use super::*;
286
287    #[cube]
288    impl<N: Size> Vector<bool, N> {
289        /// Return a new vector with the element-wise and of the vectors
290        pub fn vec_and(self, other: Self) -> Vector<bool, N> {
291            intrinsic!(
292                |scope| binary_expand(scope, self.expand, other.expand, Operator::And).into()
293            )
294        }
295    }
296}
297
298mod bool_or {
299    use cubecl_ir::Operator;
300
301    use crate::prelude::binary_expand;
302
303    use super::*;
304
305    #[cube]
306    impl<N: Size> Vector<bool, N> {
307        /// Return a new vector with the element-wise and of the vectors
308        pub fn or(self, other: Self) -> Vector<bool, N> {
309            intrinsic!(|scope| binary_expand(scope, self.expand, other.expand, Operator::Or).into())
310        }
311    }
312}
313
314impl<P: Scalar, N: Size> CubeType for Vector<P, N> {
315    type ExpandType = NativeExpand<Self>;
316}
317
318impl<P: Scalar, N: Size> CubeDebug for Vector<P, N> {}
319
320impl<P: Scalar, N: Size> NativeAssign for Vector<P, N> {
321    fn elem_init_mut(scope: &Scope, elem: Value) -> Value {
322        P::elem_init_mut(scope, elem)
323    }
324}
325
326impl<P: Scalar, N: Size> CubePrimitive for Vector<P, N> {
327    type Scalar = P;
328    type Size = N;
329    type WithScalar<S: Scalar> = Vector<S, N>;
330
331    fn __expand_as_type(scope: &Scope) -> Type {
332        Type::with_vector_size(P::__expand_as_type(scope), N::__expand_value(scope))
333    }
334
335    fn as_type_native() -> Option<Type> {
336        P::as_type_native().and_then(|ty| {
337            let vector_size = N::try_value_const()?;
338            Some(Type::with_vector_size(ty, vector_size))
339        })
340    }
341
342    fn from_const_value(value: ConstantValue) -> Self {
343        Self::new(P::from_const_value(value))
344    }
345}
346
347impl<T: Dot + Scalar, N: Size> Dot for Vector<T, N> {}
348impl<T: MulHi + Scalar, N: Size> MulHi for Vector<T, N> {}
349impl<T: FloatOps + Scalar, N: Size> FloatOps for Vector<T, N> {}
350impl<T: Hypot + Scalar, N: Size> Hypot for Vector<T, N> {}
351impl<T: Rhypot + Scalar, N: Size> Rhypot for Vector<T, N> {}
cubecl_core/frontend/container/vector/base.rs

cubecl_core/frontend/container/vector/
base.rs
