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_ir::Operation;
6use cubecl_macros::intrinsic;
7
8#[derive(Clone, Copy)]
9pub struct RuntimeCell<T: CubeType> {
10    #[allow(unused)]
11    value: T,
12}
13
14pub struct RuntimeCellExpand<T: CubeType> {
15    value: <T as cubecl::prelude::CubeType>::ExpandType,
16}
17impl<T: CubeType> Clone for RuntimeCellExpand<T> {
18    fn clone(&self) -> Self {
19        Self {
20            value: self.value.clone(),
21        }
22    }
23}
24impl<T: CubeType> cubecl::prelude::CubeType for RuntimeCell<T> {
25    type ExpandType = RuntimeCellExpand<T>;
26}
27impl<T: CubeType> cubecl::prelude::IntoMut for RuntimeCellExpand<T> {
28    fn into_mut(self, _scope: &mut cubecl::prelude::Scope) -> Self {
29        Self {
30            // We keep the same as a cell would do.
31            value: self.value.clone(),
32        }
33    }
34}
35impl<T: CubeType> cubecl::prelude::CubeDebug for RuntimeCellExpand<T> {}
36
37#[cube]
38impl<T: CubePrimitive> RuntimeCell<T> {
39    /// Create a new runtime cell with the given initial value.
40    #[allow(unused_variables)]
41    pub fn new(init: T) -> Self {
42        intrinsic!(|scope| {
43            let value = init_expand(scope, init.expand, true, Operation::Copy);
44            RuntimeCellExpand {
45                value: value.into(),
46            }
47        })
48    }
49
50    /// Store a new value in the cell.
51    #[allow(unused_variables)]
52    pub fn store(&self, value: T) {
53        intrinsic!(|scope| {
54            expand_no_check(scope, value, self.value);
55        })
56    }
57
58    /// Get the value from the call
59    pub fn read(&self) -> T {
60        intrinsic!(|scope| {
61            let value = init_expand(scope, self.value.expand, false, Operation::Copy);
62            value.into()
63        })
64    }
65
66    /// Consume the cell.
67    pub fn consume(self) -> T {
68        intrinsic!(|scope| { self.value })
69    }
70}
71
72#[cube]
73impl<T: CubeIndexMut> RuntimeCell<T> {
74    /// Store a new value in the cell at the given index.
75    #[allow(unused_variables)]
76    pub fn store_at(&mut self, index: u32, value: T::Output) {
77        intrinsic!(|scope| { self.value.expand_index_mut(scope, index, value) })
78    }
79}
80
81#[cube]
82impl<T: CubeIndex> RuntimeCell<T> {
83    /// Read a value in the cell at the given index.
84    #[allow(unused_variables)]
85    pub fn read_at(&self, index: u32) -> T::Output {
86        intrinsic!(|scope| { self.value.expand_index(scope, index) })
87    }
88}
cubecl_core/frontend/container/cell/runtime.rs

cubecl_core/frontend/container/cell/
runtime.rs
