cubecl_core/frontend/operation/
assignation.rs

1use core::ops::{
2    Index, IndexMut, Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive,
3};
4
5use cubecl_ir::{ExpandValue, Scope};
6
7use crate::unexpanded;
8use crate::{
9    frontend::{Array, Tensor},
10    prelude::*,
11};
12
13type ArrayExpand<E> = NativeExpand<Array<E>>;
14
15pub mod assign {
16    use cubecl_ir::{dialect::memory::StoreOp, interfaces::TypedExt, pliron::r#type::Typed};
17    use pliron::value::Value;
18
19    use crate::prelude::NativeExpand;
20
21    use super::*;
22
23    /// Expand the assign operation.
24    ///
25    /// If you want to assign to a manually initialized const variable, look into
26    /// [`expand_no_check()`].
27    pub fn expand<C: CubeType>(
28        scope: &Scope,
29        input: NativeExpand<C>,
30        output: &mut NativeExpand<C>,
31    ) {
32        if output.value(scope).is_immutable(scope.ctx()) {
33            panic!("Can't assign a value to a const variable. Try to use `RuntimeCell`.");
34        }
35
36        expand_no_check(scope, input, output);
37    }
38    /// Expand the assign operation without any check.
39    ///
40    /// You can't assign to a const variable with this [`expand()`].
41    pub fn expand_no_check<C: CubeType>(
42        scope: &Scope,
43        input: NativeExpand<C>,
44        output: &mut NativeExpand<C>,
45    ) {
46        let output = output.expand;
47        let input = input.expand;
48
49        expand_element(scope, input, output);
50    }
51
52    pub fn expand_element(scope: &Scope, input: ExpandValue, output: ExpandValue) {
53        let input = input.read_value(scope);
54        let output = output.value(scope);
55
56        let input = broadcast_input(scope, input, output).unwrap_or(input);
57
58        // value -> ptr = store
59        let store = StoreOp::new(scope.ctx_mut(), output, input);
60        scope.register(&store);
61    }
62
63    fn broadcast_input(scope: &Scope, input: Value, output: Value) -> Option<Value> {
64        let out_vec = output.try_get_vector_size(scope.ctx())?;
65        let in_vec = input.try_get_vector_size(scope.ctx())?;
66        if out_vec > 1 && in_vec == 1 {
67            let out_ty = output.get_type(scope.ctx()).unwrap_ptr(scope.ctx());
68            Some(cast_value(scope, input, out_ty))
69        } else {
70            None
71        }
72    }
73}
74
75pub mod index_mut {
76    use super::*;
77
78    macro_rules! impl_index {
79        ($type: ty, $expand: ty) => {
80            impl<E: CubePrimitive> IndexMut<usize> for $type {
81                fn index_mut(&mut self, _idx: usize) -> &mut Self::Output {
82                    unexpanded!()
83                }
84            }
85
86            impl<E: CubePrimitive> IndexMut<Range<usize>> for $type {
87                fn index_mut(&mut self, _idx: Range<usize>) -> &mut Self::Output {
88                    unexpanded!()
89                }
90            }
91
92            impl<E: CubePrimitive> IndexMut<RangeFrom<usize>> for $type {
93                fn index_mut(&mut self, _idx: RangeFrom<usize>) -> &mut Self::Output {
94                    unexpanded!()
95                }
96            }
97
98            impl<E: CubePrimitive> IndexMut<RangeFull> for $type {
99                fn index_mut(&mut self, _idx: RangeFull) -> &mut Self::Output {
100                    unexpanded!()
101                }
102            }
103
104            impl<E: CubePrimitive> IndexMut<RangeInclusive<usize>> for $type {
105                fn index_mut(&mut self, _idx: RangeInclusive<usize>) -> &mut Self::Output {
106                    unexpanded!()
107                }
108            }
109
110            impl<E: CubePrimitive> IndexMut<RangeTo<usize>> for $type {
111                fn index_mut(&mut self, _idx: RangeTo<usize>) -> &mut Self::Output {
112                    unexpanded!()
113                }
114            }
115
116            impl<E: CubePrimitive> IndexMut<RangeToInclusive<usize>> for $type {
117                fn index_mut(&mut self, _idx: RangeToInclusive<usize>) -> &mut Self::Output {
118                    unexpanded!()
119                }
120            }
121
122            impl<E: CubePrimitive> IndexMutExpand<NativeExpand<usize>> for $expand {
123                fn __expand_index_mut_method(
124                    &mut self,
125                    scope: &Scope,
126                    index: NativeExpand<usize>,
127                ) -> &mut E::ExpandType {
128                    self.__expand_as_mut_slice_method(scope)
129                        .__expand_index_mut_method(scope, index)
130                }
131            }
132        };
133    }
134
135    impl<E: CubePrimitive, const N: usize> IndexMutExpand<NativeExpand<usize>>
136        for NativeExpand<[E; N]>
137    {
138        fn __expand_index_mut_method(
139            &mut self,
140            scope: &Scope,
141            index: NativeExpand<usize>,
142        ) -> &mut E::ExpandType {
143            self.__expand_as_mut_slice_method(scope)
144                .__expand_index_mut_method(scope, index)
145        }
146    }
147
148    impl_index!(Array<E>, NativeExpand<Array<E>>);
149    impl_index!(Tensor<E>, TensorExpand<E>);
150    impl_index!(Shared<[E]>, NativeExpand<Shared<[E]>>);
151
152    impl_slice_ranges!([E: CubePrimitive] ArrayExpand<E>);
153    impl_slice_ranges!([E: CubePrimitive, const N: usize] NativeExpand<[E; N]>);
154    impl_slice_ranges!([E: CubePrimitive] TensorExpand<E>);
155    impl_slice_ranges!([E: CubePrimitive] NativeExpand<Shared<[E]>>);
156}
157
158pub mod index {
159    use super::*;
160
161    macro_rules! impl_index {
162        ($type: ty, $expand: ty) => {
163            impl<E: CubePrimitive> Index<usize> for $type {
164                type Output = E;
165
166                fn index(&self, _idx: usize) -> &Self::Output {
167                    unexpanded!()
168                }
169            }
170
171            impl<E: CubePrimitive> Index<Range<usize>> for $type {
172                type Output = [E];
173
174                fn index(&self, _idx: Range<usize>) -> &Self::Output {
175                    unexpanded!()
176                }
177            }
178
179            impl<E: CubePrimitive> Index<RangeFrom<usize>> for $type {
180                type Output = [E];
181
182                fn index(&self, _idx: RangeFrom<usize>) -> &Self::Output {
183                    unexpanded!()
184                }
185            }
186
187            impl<E: CubePrimitive> Index<RangeFull> for $type {
188                type Output = [E];
189
190                fn index(&self, _idx: RangeFull) -> &Self::Output {
191                    unexpanded!()
192                }
193            }
194
195            impl<E: CubePrimitive> Index<RangeInclusive<usize>> for $type {
196                type Output = [E];
197
198                fn index(&self, _idx: RangeInclusive<usize>) -> &Self::Output {
199                    unexpanded!()
200                }
201            }
202
203            impl<E: CubePrimitive> Index<RangeTo<usize>> for $type {
204                type Output = [E];
205
206                fn index(&self, _idx: RangeTo<usize>) -> &Self::Output {
207                    unexpanded!()
208                }
209            }
210
211            impl<E: CubePrimitive> Index<RangeToInclusive<usize>> for $type {
212                type Output = [E];
213
214                fn index(&self, _idx: RangeToInclusive<usize>) -> &Self::Output {
215                    unexpanded!()
216                }
217            }
218
219            impl<E: CubePrimitive> IndexExpand<NativeExpand<usize>> for $expand {
220                type Output = NativeExpand<E>;
221
222                fn __expand_index_method(
223                    &self,
224                    scope: &Scope,
225                    index: NativeExpand<usize>,
226                ) -> &Self::Output {
227                    self.__expand_as_slice_method(scope)
228                        .__expand_index_method(scope, index)
229                }
230            }
231        };
232    }
233
234    impl<E: CubePrimitive, const N: usize> IndexExpand<NativeExpand<usize>> for NativeExpand<[E; N]> {
235        type Output = NativeExpand<E>;
236
237        fn __expand_index_method(
238            &self,
239            scope: &Scope,
240            index: NativeExpand<usize>,
241        ) -> &Self::Output {
242            self.__expand_as_slice_method(scope)
243                .__expand_index_method(scope, index)
244        }
245    }
246
247    impl_index!(Array<E>, NativeExpand<Array<E>>);
248    impl_index!(Tensor<E>, TensorExpand<E>);
249    impl_index!(Shared<[E]>, NativeExpand<Shared<[E]>>);
250}
cubecl_core/frontend/operation/assignation.rs

cubecl_core/frontend/operation/
assignation.rs
