cubecl_core/frontend/container/
shared_memory.rs

1use core::ops::{Deref, DerefMut};
2
3use crate::{
4    self as cubecl,
5    frontend::container::slice,
6    prelude::{Vectorized, VectorizedExpand},
7    unexpanded,
8};
9use cubecl_ir::{VectorSize, dialect::general::FreeOp, interfaces::TypedExt, types::ArrayType};
10use cubecl_macros::{cube, intrinsic};
11
12use crate::{
13    frontend::{CubePrimitive, CubeType, IntoMut, NativeExpand},
14    ir::Scope,
15    prelude::*,
16};
17
18pub type SharedExpand<T> = NativeExpand<Shared<T>>;
19
20pub struct Shared<E: NativeCubeType + ?Sized> {
21    _val: *mut E,
22}
23
24// Treat it as a shared smart pointer
25impl<E: NativeCubeType + ?Sized> Clone for Shared<E> {
26    fn clone(&self) -> Self {
27        Self { _val: self._val }
28    }
29}
30
31impl<T: NativeCubeType + ?Sized> IntoMut for NativeExpand<Shared<T>> {
32    fn into_mut(self, _scope: &Scope) -> Self {
33        self
34    }
35}
36
37impl<T: NativeCubeType + ?Sized> CubeType for Shared<T> {
38    type ExpandType = NativeExpand<Shared<T>>;
39}
40
41impl<T: NativeCubeType + ?Sized> AsMutExpand for SharedExpand<T> {
42    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
43        self
44    }
45}
46
47#[cube]
48impl<T: CubePrimitive> Shared<[T]> {
49    /// Create a new shared slice.
50    ///
51    /// # Safety
52    /// Shared memory is always uninitialized by default. Reading uninitialized shared values is
53    /// undefined behavior.
54    pub fn new_slice(#[comptime] len: usize) -> Self {
55        intrinsic!(|scope| {
56            let inner = T::__expand_as_type(scope);
57            let ty = ArrayType::get(scope.ctx(), inner, len);
58            let buffer = scope.create_shared(ty, None);
59            let slice = slice::from_raw_parts::<T>(
60                scope,
61                buffer,
62                0usize.into_expand(scope),
63                len.into_expand(scope),
64            );
65            slice.expand.into()
66        })
67    }
68
69    /// Create a new shared slice with a specified minimum alignment.
70    ///
71    /// # Safety
72    /// Shared memory is always uninitialized by default. Reading uninitialized shared values is
73    /// undefined behavior.
74    #[allow(unused_variables)]
75    pub fn new_aligned_slice(#[comptime] len: usize, #[comptime] alignment: usize) -> Self {
76        intrinsic!(|scope| {
77            let inner = T::__expand_as_type(scope);
78            let ty = ArrayType::get(scope.ctx(), inner, len);
79            let buffer = scope.create_shared(ty, Some(alignment));
80            let slice = slice::from_raw_parts::<T>(
81                scope,
82                buffer,
83                0usize.into_expand(scope),
84                len.into_expand(scope),
85            );
86            slice.expand.into()
87        })
88    }
89
90    #[allow(clippy::len_without_is_empty)]
91    pub fn len(&self) -> usize {
92        intrinsic!(|scope| len_static(scope, &self))
93    }
94}
95
96impl<T: NativeCubeType + ?Sized> Shared<T> {
97    pub fn map<U: NativeCubeType + ?Sized>(self, _map: impl FnOnce(&T) -> &U) -> Shared<U> {
98        unexpanded!()
99    }
100}
101
102impl<T: NativeCubeType + ?Sized> SharedExpand<T> {
103    pub fn __expand_map_method<U: NativeCubeType + ?Sized>(
104        self,
105        scope: &Scope,
106        map: impl for<'a> FnOnce(&Scope, &'a NativeExpand<T>) -> &'a NativeExpand<U>,
107    ) -> SharedExpand<U> {
108        let out = map(scope, self.__expand_ref_method(scope));
109        out.expand.into()
110    }
111}
112
113#[cube]
114impl<T: CubePrimitive> Shared<T> {
115    /// Create a new shared object.
116    ///
117    /// # Safety
118    /// Shared memory is always uninitialized by default. Reading uninitialized shared values is
119    /// undefined behavior.
120    pub fn new() -> Self {
121        intrinsic!(|scope| {
122            let val = scope.create_shared(T::__expand_as_type(scope), None);
123            NativeExpand::new(val.into())
124        })
125    }
126}
127
128#[cube]
129impl<T: NativeCubeType + ?Sized> Shared<T> {
130    pub fn inner_ref(&self) -> &T {
131        intrinsic!(|scope| { unsafe { self.as_type_ref_unchecked() } })
132    }
133
134    pub fn inner_mut(&mut self) -> &mut T {
135        intrinsic!(|scope| { unsafe { self.as_type_mut_unchecked() } })
136    }
137}
138
139impl<T: CubePrimitive> Default for Shared<T> {
140    fn default() -> Self {
141        Self::new()
142    }
143}
144impl<T: CubePrimitive> Shared<T> {
145    pub fn __expand_default(scope: &Scope) -> <Self as CubeType>::ExpandType {
146        Self::__expand_new(scope)
147    }
148}
149
150#[cube]
151impl<T: NativeCubeType + ?Sized> Shared<T> {
152    /// Frees the shared memory for reuse, if possible on the target runtime.
153    ///
154    /// # Safety
155    /// *Must* be used in uniform control flow
156    /// *Must not* have any dangling references to this shared memory
157    pub unsafe fn free(&self) {
158        intrinsic!(|scope| {
159            let val = scope.extract_field(self.value(scope), 0);
160            scope.register(&FreeOp::new(scope.ctx_mut(), val))
161        })
162    }
163}
164
165fn len_static<T: CubePrimitive>(
166    scope: &Scope,
167    shared: &NativeExpand<Shared<[T]>>,
168) -> NativeExpand<usize> {
169    let array_ty = inner_array_ty(scope, shared.value(scope));
170    array_ty.deref(scope.ctx()).length.into()
171}
172
173impl<T: CubePrimitive> List<T> for Shared<[T]> {}
174impl<T: CubePrimitive> ListExpand<T> for NativeExpand<Shared<[T]>> {
175    fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
176        Self::__expand_len_method(self, scope)
177    }
178}
179
180impl<T: CubePrimitive> Vectorized for Shared<[T]> {}
181impl<T: CubePrimitive> VectorizedExpand for NativeExpand<Shared<[T]>> {
182    fn __expand_vector_size_method(&self, scope: &Scope) -> VectorSize {
183        self.__extract_list(scope).vector_size(scope.ctx())
184    }
185}
186
187impl<T: NativeCubeType + ?Sized> Deref for Shared<T> {
188    type Target = T;
189
190    fn deref(&self) -> &Self::Target {
191        unexpanded!()
192    }
193}
194impl<T: NativeCubeType + ?Sized> DerefMut for Shared<T> {
195    fn deref_mut(&mut self) -> &mut Self::Target {
196        unexpanded!()
197    }
198}
199
200impl<T: NativeCubeType + ?Sized> Deref for SharedExpand<T> {
201    type Target = NativeExpand<T>;
202
203    fn deref(&self) -> &Self::Target {
204        unsafe { self.as_type_ref_unchecked() }
205    }
206}
207impl<T: NativeCubeType + ?Sized> DerefMut for SharedExpand<T> {
208    fn deref_mut(&mut self) -> &mut Self::Target {
209        unsafe { self.as_type_mut_unchecked() }
210    }
211}
212
213impl<'a, T: NativeCubeType + ?Sized> From<&'a SharedExpand<T>> for &'a NativeExpand<T> {
214    fn from(value: &'a SharedExpand<T>) -> Self {
215        value
216    }
217}
218
219impl<'a, T: NativeCubeType + ?Sized> From<&'a mut SharedExpand<T>> for &'a mut NativeExpand<T> {
220    fn from(value: &'a mut SharedExpand<T>) -> Self {
221        value
222    }
223}
224
225impl<T: NativeCubeType + ?Sized> AsDerefExpand for SharedExpand<T> {
226    type Target = NativeExpand<T>;
227
228    fn __expand_as_deref_method(&self, _: &Scope) -> &Self::Target {
229        unsafe { self.as_type_ref_unchecked::<T>() }
230    }
231}
232impl<T: NativeCubeType + ?Sized> AsDerefMutExpand for SharedExpand<T> {
233    fn __expand_as_deref_mut_method(&mut self, _: &Scope) -> &mut Self::Target {
234        unsafe { self.as_type_mut_unchecked::<T>() }
235    }
236}
237
238impl<T: CubePrimitive> Assign<NativeExpand<T>> for SharedExpand<T> {
239    fn __expand_assign_method(&mut self, scope: &Scope, value: NativeExpand<T>) {
240        let value = value.read_value(scope);
241        assign::expand_element(scope, value.into(), self.expand);
242    }
243}
244
245impl<T: CubePrimitive> RuntimeAssign<NativeExpand<T>> for SharedExpand<T> {
246    fn init_mut(&self, scope: &Scope) -> Self::Expand {
247        Shared::<T>::__expand_new(scope)
248    }
249}
cubecl_core/frontend/container/shared_memory.rs

cubecl_core/frontend/container/
shared_memory.rs
