Skip to main content

cubecl_core/frontend/
list.rs

1use super::CubeType;
2use crate as cubecl;
3use crate::{prelude::*, unexpanded};
4use cubecl_ir::{Scope, VectorSize};
5
6/// Type from which we can read/to which we can write values in cube functions.
7#[allow(clippy::len_without_is_empty)]
8#[cube(expand_base_traits = "SliceOperatorExpand<T>
9    + IndexExpand<NativeExpand<usize>, Output = NativeExpand<T>>
10    + IndexMutExpand<NativeExpand<usize>, Output = NativeExpand<T>>")]
11pub trait List<T: CubePrimitive>:
12    SliceOperator<T> + CubeIndex<usize, Output = T> + CubeIndexMut<usize, Output = T> + Vectorized
13{
14    fn len(&self) -> usize {
15        unexpanded!();
16    }
17}
18
19pub trait Vectorized: CubeType<ExpandType: VectorizedExpand> {
20    fn vector_size(&self) -> VectorSize {
21        unexpanded!()
22    }
23    fn __expand_vector_size(_scope: &Scope, this: &Self::ExpandType) -> VectorSize {
24        this.vector_size()
25    }
26}
27
28impl<T: Vectorized + ?Sized> Vectorized for &T {}
29impl<T: Vectorized + ?Sized> Vectorized for &mut T {}
30
31pub trait VectorizedExpand {
32    fn vector_size(&self) -> VectorSize;
33    fn __expand_vector_size_method(&self, _scope: &Scope) -> VectorSize {
34        self.vector_size()
35    }
36}
37
38impl<T: VectorizedExpand + ?Sized> VectorizedExpand for &T {
39    fn vector_size(&self) -> VectorSize {
40        (**self).vector_size()
41    }
42}
43impl<T: VectorizedExpand + ?Sized> VectorizedExpand for &mut T {
44    fn vector_size(&self) -> VectorSize {
45        (**self).vector_size()
46    }
47}