cubecl_core/frontend/
list.rs

1use super::{CubeType, ExpandElementTyped};
2use crate as cubecl;
3use crate::{prelude::*, unexpanded};
4use cubecl_ir::Scope;
5
6/// Type from which we can read values in cube functions.
7/// For a mutable version, see [ListMut].
8#[allow(clippy::len_without_is_empty)]
9#[cube(self_type = "ref", expand_base_traits = "SliceOperatorExpand<T>")]
10pub trait List<T: CubePrimitive>: SliceOperator<T> + Lined {
11    #[allow(unused)]
12    fn read(&self, index: u32) -> T {
13        unexpanded!()
14    }
15
16    #[allow(unused)]
17    fn read_unchecked(&self, index: u32) -> T {
18        unexpanded!()
19    }
20
21    #[allow(unused)]
22    fn len(&self) -> u32 {
23        unexpanded!();
24    }
25}
26
27/// Type for which we can read and write values in cube functions.
28/// For an immutable version, see [List].
29#[cube(self_type = "ref", expand_base_traits = "SliceMutOperatorExpand<T>")]
30pub trait ListMut<T: CubePrimitive>: List<T> + SliceMutOperator<T> {
31    #[allow(unused)]
32    fn write(&self, index: u32, value: T) {
33        unexpanded!()
34    }
35}
36
37// Automatic implementation for references to List.
38impl<'a, T: CubePrimitive, L: List<T>> List<T> for &'a L
39where
40    &'a L: CubeType<ExpandType = L::ExpandType>,
41{
42    fn read(&self, index: u32) -> T {
43        L::read(self, index)
44    }
45
46    fn __expand_read(
47        scope: &mut Scope,
48        this: Self::ExpandType,
49        index: ExpandElementTyped<u32>,
50    ) -> <T as CubeType>::ExpandType {
51        L::__expand_read(scope, this, index)
52    }
53}
54
55// Automatic implementation for mutable references to List.
56impl<'a, T: CubePrimitive, L: List<T>> List<T> for &'a mut L
57where
58    &'a mut L: CubeType<ExpandType = L::ExpandType>,
59{
60    fn read(&self, index: u32) -> T {
61        L::read(self, index)
62    }
63
64    fn __expand_read(
65        scope: &mut Scope,
66        this: Self::ExpandType,
67        index: ExpandElementTyped<u32>,
68    ) -> <T as CubeType>::ExpandType {
69        L::__expand_read(scope, this, index)
70    }
71}
72
73// Automatic implementation for references to ListMut.
74impl<'a, T: CubePrimitive, L: ListMut<T>> ListMut<T> for &'a L
75where
76    &'a L: CubeType<ExpandType = L::ExpandType>,
77{
78    fn write(&self, index: u32, value: T) {
79        L::write(self, index, value);
80    }
81
82    fn __expand_write(
83        scope: &mut Scope,
84        this: Self::ExpandType,
85        index: ExpandElementTyped<u32>,
86        value: T::ExpandType,
87    ) {
88        L::__expand_write(scope, this, index, value);
89    }
90}
91
92// Automatic implementation for mutable references to ListMut.
93impl<'a, T: CubePrimitive, L: ListMut<T>> ListMut<T> for &'a mut L
94where
95    &'a mut L: CubeType<ExpandType = L::ExpandType>,
96{
97    fn write(&self, index: u32, value: T) {
98        L::write(self, index, value);
99    }
100
101    fn __expand_write(
102        scope: &mut Scope,
103        this: Self::ExpandType,
104        index: ExpandElementTyped<u32>,
105        value: T::ExpandType,
106    ) {
107        L::__expand_write(scope, this, index, value);
108    }
109}
110
111pub trait Lined: CubeType<ExpandType: LinedExpand> {
112    fn line_size(&self) -> u32 {
113        unexpanded!()
114    }
115    fn __expand_line_size(_scope: &mut Scope, this: Self::ExpandType) -> u32 {
116        this.line_size()
117    }
118}
119
120pub trait LinedExpand {
121    fn line_size(&self) -> u32;
122    fn __expand_line_size_method(&self, _scope: &mut Scope) -> u32 {
123        self.line_size()
124    }
125}
126
127impl<'a, L: Lined> Lined for &'a L where &'a L: CubeType<ExpandType: LinedExpand> {}
128impl<'a, L: Lined> Lined for &'a mut L where &'a mut L: CubeType<ExpandType: LinedExpand> {}
cubecl_core/frontend/list.rs

cubecl_core/frontend/
list.rs
