cubecl_core/frontend/container/sequence/
base.rs

1use serde::{Deserialize, Serialize};
2
3use crate::frontend::{
4    branch::Iterable, indexation::Index, CubeContext, CubeType, ExpandElementTyped, Init,
5    IntoRuntime,
6};
7use crate::prelude::ExpandElement;
8use std::{cell::RefCell, rc::Rc};
9
10/// A sequence of [cube types](CubeType) that is inlined during compilation.
11///
12/// In other words, it allows you to group a dynamic amount of variables at compile time.
13///
14/// All methods [push](Sequence::push), [index](Sequence::index) and
15/// [into_iter](Sequence::into_iter) are executed _during_ compilation and don't add any overhead
16/// on the generated kernel.
17#[derive(Debug, Clone, Hash, PartialEq, Eq, Serialize, Deserialize)]
18pub struct Sequence<T: CubeType> {
19    values: Vec<T>,
20}
21
22impl<T: CubeType> Default for Sequence<T> {
23    fn default() -> Self {
24        Self::new()
25    }
26}
27
28impl<T: CubeType> Init for Sequence<T> {
29    fn init(self, _context: &mut CubeContext) -> Self {
30        self
31    }
32}
33
34impl<T: CubeType> Sequence<T> {
35    /// Create a new empty sequence.
36    pub fn new() -> Self {
37        Self { values: Vec::new() }
38    }
39
40    /// Push a new value into the sequence.
41    pub fn push(&mut self, value: T) {
42        self.values.push(value);
43    }
44
45    /// Obtain the sequence length.
46    #[allow(clippy::len_without_is_empty)]
47    pub fn len(&self) -> u32 {
48        self.values.len() as u32
49    }
50
51    /// Get the variable at the given position in the sequence.
52    #[allow(unused_variables, clippy::should_implement_trait)]
53    pub fn index<I: Index>(&self, index: I) -> &T {
54        let index: ExpandElementTyped<u32> = ExpandElement::Plain(index.value()).into();
55        let index = index
56            .constant()
57            .expect("Only constant are supported")
58            .as_usize();
59
60        self.values.get(index).unwrap()
61    }
62
63    /// Get the variable at the given position in the sequence.
64    #[allow(unused_variables, clippy::should_implement_trait)]
65    pub fn index_mut<I: Index>(&mut self, index: I) -> &mut T {
66        let index: ExpandElementTyped<u32> = ExpandElement::Plain(index.value()).into();
67        let index = index
68            .constant()
69            .expect("Only constant are supported")
70            .as_usize();
71
72        self.values.get_mut(index).unwrap()
73    }
74
75    /// Expand function of [new](Self::new).
76    pub fn __expand_new(_context: &mut CubeContext) -> SequenceExpand<T> {
77        SequenceExpand {
78            values: Rc::new(RefCell::new(Vec::new())),
79        }
80    }
81
82    /// Insert an item at the given index.
83    #[allow(unused_variables, clippy::should_implement_trait)]
84    pub fn insert<I: Index>(&mut self, index: I, value: T) {
85        *self.index_mut(index) = value;
86    }
87
88    /// Expand function of [push](Self::push).
89    pub fn __expand_push(
90        context: &mut CubeContext,
91        expand: &mut SequenceExpand<T>,
92        value: T::ExpandType,
93    ) {
94        expand.__expand_push_method(context, value)
95    }
96
97    /// Expand function of [index](Self::index).
98    pub fn __expand_index(
99        context: &mut CubeContext,
100        expand: SequenceExpand<T>,
101        index: ExpandElementTyped<u32>,
102    ) -> T::ExpandType {
103        expand.__expand_index_method(context, index)
104    }
105
106    /// Expand function of [index_mut](Self::index_mut).
107    pub fn __expand_index_mut(
108        context: &mut CubeContext,
109        expand: SequenceExpand<T>,
110        index: ExpandElementTyped<u32>,
111    ) -> T::ExpandType {
112        expand.__expand_index_mut_method(context, index)
113    }
114}
115
116/// Expand type of [Sequence].
117pub struct SequenceExpand<T: CubeType> {
118    // We clone the expand type during the compilation phase, but for register reuse, not for
119    // copying data. To achieve the intended behavior, we have to share the same underlying values.
120    pub(super) values: Rc<RefCell<Vec<T::ExpandType>>>,
121}
122
123impl<T: CubeType> Iterable<T> for SequenceExpand<T> {
124    fn expand(
125        self,
126        context: &mut CubeContext,
127        func: impl FnMut(&mut CubeContext, <T as CubeType>::ExpandType),
128    ) {
129        self.expand_unroll(context, func);
130    }
131
132    fn expand_unroll(
133        self,
134        context: &mut CubeContext,
135        mut func: impl FnMut(&mut CubeContext, <T as CubeType>::ExpandType),
136    ) {
137        for elem in self {
138            func(context, elem);
139        }
140    }
141}
142
143impl<T: CubeType> Init for SequenceExpand<T> {
144    fn init(self, _context: &mut crate::prelude::CubeContext) -> Self {
145        self
146    }
147}
148
149impl<T: CubeType> Clone for SequenceExpand<T> {
150    fn clone(&self) -> Self {
151        Self {
152            values: self.values.clone(),
153        }
154    }
155}
156
157impl<T: CubeType> IntoIterator for Sequence<T> {
158    type Item = T;
159
160    type IntoIter = <Vec<T> as IntoIterator>::IntoIter;
161
162    fn into_iter(self) -> Self::IntoIter {
163        self.values.into_iter()
164    }
165}
166
167impl<T: CubeType> IntoIterator for SequenceExpand<T> {
168    type Item = T::ExpandType;
169
170    type IntoIter = <Vec<T::ExpandType> as IntoIterator>::IntoIter;
171
172    fn into_iter(self) -> Self::IntoIter {
173        self.values.take().into_iter()
174    }
175}
176
177impl<T: CubeType> CubeType for Sequence<T> {
178    type ExpandType = SequenceExpand<T>;
179}
180
181impl<T: CubeType> SequenceExpand<T> {
182    /// Expand method of [push](Sequence::push).
183    pub fn __expand_push_method(&mut self, _context: &mut CubeContext, value: T::ExpandType) {
184        self.values.borrow_mut().push(value);
185    }
186
187    /// Expand method of [insert](Sequence::insert).
188    pub fn __expand_insert_method(
189        &self,
190        _context: &mut CubeContext,
191        index: ExpandElementTyped<u32>,
192        value: T::ExpandType,
193    ) {
194        let index = index
195            .constant()
196            .expect("Only constant are supported")
197            .as_usize();
198
199        let mut values = self.values.borrow_mut();
200
201        if values.len() == index {
202            values.push(value);
203        } else {
204            values[index] = value;
205        }
206    }
207
208    /// Expand method of [index](Sequence::index).
209    pub fn __expand_index_method(
210        &self,
211        _context: &mut CubeContext,
212        index: ExpandElementTyped<u32>,
213    ) -> T::ExpandType {
214        let index = index
215            .constant()
216            .expect("Only constant are supported")
217            .as_usize();
218
219        self.values.borrow()[index].clone()
220    }
221
222    /// Expand method of [index_mut](Sequence::index_mut).
223    pub fn __expand_index_mut_method(
224        &self,
225        _context: &mut CubeContext,
226        index: ExpandElementTyped<u32>,
227    ) -> T::ExpandType {
228        let index = index
229            .constant()
230            .expect("Only constant are supported")
231            .as_usize();
232
233        self.values.borrow()[index].clone()
234    }
235
236    pub fn __expand_len_method(&self, _context: &mut CubeContext) -> u32 {
237        let values = self.values.borrow();
238        values.len() as u32
239    }
240}
241
242impl<T: CubeType> IntoRuntime for Sequence<T> {
243    fn __expand_runtime_method(self, _context: &mut CubeContext) -> SequenceExpand<T> {
244        unimplemented!("Sequence doesn't exist at compile time");
245    }
246}
cubecl_core/frontend/container/sequence/base.rs

cubecl_core/frontend/container/sequence/
base.rs
