cubecl_core/frontend/container/vector/
base.rs

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

cubecl_core/frontend/container/vector/
base.rs
