cubecl_core/frontend/container/cell/
comptime.rs

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

cubecl_core/frontend/container/cell/
comptime.rs
