cubecl_core/frontend/container/cell/
comptime.rs

1use alloc::rc::Rc;
2use core::cell::RefCell;
3
4use cubecl_ir::Scope;
5
6use crate::prelude::{CubeDebug, CubeType, IntoMut};
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> CubeType for ComptimeCell<T> {
22    type ExpandType = ComptimeCellExpand<T>;
23}
24
25impl<T: CubeType + Clone> ComptimeCell<T> {
26    pub fn new(value: T) -> Self {
27        Self {
28            value: Rc::new(RefCell::new(value)),
29        }
30    }
31    pub fn __expand_new(_scope: &mut Scope, value: T::ExpandType) -> ComptimeCellExpand<T> {
32        ComptimeCellExpand {
33            value: Rc::new(RefCell::new(value)),
34        }
35    }
36    pub fn read(&self) -> T {
37        let value = self.value.borrow();
38        value.clone()
39    }
40    pub fn store(&mut self, value: T) {
41        let mut old = self.value.borrow_mut();
42        *old = value;
43    }
44    pub fn __expand_store(context: &mut Scope, this: ComptimeCellExpand<T>, value: T::ExpandType) {
45        this.__expand_store_method(context, value)
46    }
47    pub fn __expand_read(scope: &mut Scope, this: ComptimeCellExpand<T>) -> T::ExpandType {
48        this.__expand_read_method(scope)
49    }
50}
51
52impl<T: CubeType + Clone> ComptimeCellExpand<T> {
53    pub fn __expand_store_method(&self, _context: &mut Scope, value: T::ExpandType) {
54        let mut old = self.value.borrow_mut();
55        *old = value;
56    }
57    pub fn __expand_read_method(&self, _scope: &mut Scope) -> T::ExpandType {
58        let value = self.value.borrow();
59        value.clone()
60    }
61}
62
63impl<T: CubeType> IntoMut for ComptimeCellExpand<T> {
64    fn into_mut(self, _scope: &mut Scope) -> Self {
65        self
66    }
67}
68
69impl<T: CubeType> CubeDebug for ComptimeCellExpand<T> {}
70
71impl<T: CubeType> Clone for ComptimeCellExpand<T> {
72    fn clone(&self) -> Self {
73        Self {
74            value: self.value.clone(),
75        }
76    }
77}
cubecl_core/frontend/container/cell/comptime.rs

cubecl_core/frontend/container/cell/
comptime.rs
