cubecl_core/frontend/container/cell/
runtime.rs

1use crate as cubecl;
2use crate::prelude::CubePrimitive;
3use cubecl::frontend::assign::expand_no_check;
4use cubecl::prelude::*;
5use cubecl_macros::intrinsic;
6
7#[derive(Clone, Copy)]
8pub struct RuntimeCell<T: CubeType> {
9    #[allow(unused)]
10    value: T,
11}
12
13pub struct RuntimeCellExpand<T: CubeType> {
14    value: <T as cubecl::prelude::CubeType>::ExpandType,
15}
16impl<T: CubeType<ExpandType: Clone>> Clone for RuntimeCellExpand<T> {
17    fn clone(&self) -> Self {
18        Self {
19            value: self.value.clone(),
20        }
21    }
22}
23
24impl<T: CubeType> IntoExpand for RuntimeCellExpand<T> {
25    type Expand = Self;
26
27    fn into_expand(self, _scope: &Scope) -> Self::Expand {
28        self
29    }
30}
31impl<T: CubeType> ExpandTypeClone for RuntimeCellExpand<T> {
32    fn clone_unchecked(&self) -> Self {
33        Self {
34            value: self.value.clone_unchecked(),
35        }
36    }
37}
38impl<T: CubeType> cubecl::prelude::CubeType for RuntimeCell<T> {
39    type ExpandType = RuntimeCellExpand<T>;
40}
41
42impl<T: CubeType> cubecl::prelude::IntoMut for RuntimeCellExpand<T> {
43    fn into_mut(self, _scope: &Scope) -> Self {
44        Self {
45            // We keep the same as a cell would do.
46            value: self.value,
47        }
48    }
49}
50impl<T: CubeType> cubecl::prelude::CubeDebug for RuntimeCellExpand<T> {}
51
52#[cube]
53impl<T: CubePrimitive> RuntimeCell<T> {
54    /// Create a new runtime cell with the given initial value.
55    pub fn new(init: T) -> Self {
56        intrinsic!(|scope| {
57            let value = init_expand(scope, init.expand, true);
58            RuntimeCellExpand {
59                value: value.into(),
60            }
61        })
62    }
63
64    /// Store a new value in the cell.
65    pub fn store(&self, value: T) {
66        intrinsic!(|scope| {
67            let mut this = self.value.clone();
68            expand_no_check(scope, value, &mut this);
69        })
70    }
71
72    /// Get the value from the call
73    pub fn read(&self) -> T {
74        intrinsic!(|scope| {
75            let value = init_expand(scope, self.value.clone().expand, false);
76            value.into()
77        })
78    }
79
80    /// Consume the cell.
81    pub fn consume(self) -> T {
82        intrinsic!(|scope| { self.value })
83    }
84}
85
86#[cube]
87impl<S: Scalar, N: Size> RuntimeCell<Vector<S, N>> {
88    /// Extract the value in the cell at the given index.
89    pub fn extract(&mut self, index: usize) -> S {
90        intrinsic!(|scope| { self.value.__expand_extract_method(scope, index) })
91    }
92
93    /// Store a new value in the cell at the given index.
94    pub fn insert(&mut self, index: usize, value: S) {
95        intrinsic!(|scope| { self.value.__expand_insert_method(scope, index, value) })
96    }
97}
98
99impl<T: CubeType> AsRefExpand for RuntimeCellExpand<T> {
100    fn __expand_ref_method(&self, _: &Scope) -> &Self {
101        self
102    }
103}
104impl<T: CubeType> AsMutExpand for RuntimeCellExpand<T> {
105    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
106        self
107    }
108}
cubecl_core/frontend/container/cell/runtime.rs

cubecl_core/frontend/container/cell/
runtime.rs
