cubecl_core/frontend/container/cell/
comptime.rs

1use alloc::rc::Rc;
2use core::cell::RefCell;
3
4use cubecl_ir::Scope;
5
6use crate::prelude::*;
7
8#[derive(Debug, Clone)]
9/// A cell that can store and mutate a cube type during comptime.
10pub struct ComptimeCell<T: CubeType> {
11    pub(super) value: Rc<RefCell<T>>,
12}
13
14/// Expand type of [`ComptimeCell`].
15pub struct ComptimeCellExpand<T: CubeType> {
16    // We clone the expand type during the compilation phase, but for register reuse, not for
17    // copying data. To achieve the intended behavior, we have to share the same underlying values.
18    pub(super) value: Rc<RefCell<T::ExpandType>>,
19}
20
21impl<T: CubeType> IntoExpand for ComptimeCellExpand<T> {
22    type Expand = Self;
23
24    fn into_expand(self, _scope: &Scope) -> Self::Expand {
25        self
26    }
27}
28impl<T: CubeType> ExpandTypeClone for ComptimeCellExpand<T> {
29    fn clone_unchecked(&self) -> Self {
30        Self {
31            value: self.value.clone(),
32        }
33    }
34}
35
36impl<T: CubeType> CubeType for ComptimeCell<T> {
37    type ExpandType = ComptimeCellExpand<T>;
38}
39
40impl<T: CubeType<ExpandType: Clone> + Clone> ComptimeCell<T> {
41    pub fn new(value: T) -> Self {
42        Self {
43            value: Rc::new(RefCell::new(value)),
44        }
45    }
46    pub fn __expand_new(_scope: &Scope, value: T::ExpandType) -> ComptimeCellExpand<T> {
47        ComptimeCellExpand {
48            value: Rc::new(RefCell::new(value)),
49        }
50    }
51    pub fn read(&self) -> T {
52        let value = self.value.borrow();
53        value.clone()
54    }
55    pub fn store(&mut self, value: T) {
56        let mut old = self.value.borrow_mut();
57        *old = value;
58    }
59    pub fn __expand_store(context: &Scope, this: ComptimeCellExpand<T>, value: T::ExpandType) {
60        this.__expand_store_method(context, value)
61    }
62    pub fn __expand_read(scope: &Scope, this: ComptimeCellExpand<T>) -> T::ExpandType {
63        this.__expand_read_method(scope)
64    }
65}
66
67impl<T: CubeType<ExpandType: Clone> + Clone> ComptimeCellExpand<T> {
68    pub fn __expand_store_method(&self, _context: &Scope, value: T::ExpandType) {
69        let mut old = self.value.borrow_mut();
70        *old = value;
71    }
72    pub fn __expand_read_method(&self, _scope: &Scope) -> T::ExpandType {
73        let value: core::cell::Ref<'_, T::ExpandType> = self.value.borrow();
74        value.clone()
75    }
76}
77
78impl<T: CubeType> IntoMut for ComptimeCellExpand<T> {
79    fn into_mut(self, _scope: &Scope) -> Self {
80        self
81    }
82}
83
84impl<T: CubeType> CubeDebug for ComptimeCellExpand<T> {}
85
86impl<T: CubeType> Clone for ComptimeCellExpand<T> {
87    fn clone(&self) -> Self {
88        Self {
89            value: self.value.clone(),
90        }
91    }
92}
93
94impl<T: CubeType> AsRefExpand for ComptimeCellExpand<T> {
95    fn __expand_ref_method(&self, _: &Scope) -> &Self {
96        self
97    }
98}
99impl<T: CubeType> AsMutExpand for ComptimeCellExpand<T> {
100    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
101        self
102    }
103}
cubecl_core/frontend/container/cell/comptime.rs

cubecl_core/frontend/container/cell/
comptime.rs
