cubecl_core/frontend/
indexation.rs

1use core::ops::{Index, IndexMut};
2
3use cubecl_ir::{ExpandValue, Scope, dialect::memory::IndexOp, pliron::value::Value};
4
5use crate::frontend::ReadValue;
6
7use super::{CubeType, NativeExpand, index_expand};
8
9/// Trait bound that can be used to guarantee the expand also implements `IndexExpand`
10pub trait CubeIndex<I: CubeType>:
11    Index<I, Output: CubeType>
12    + CubeType<
13        ExpandType: IndexExpand<I::ExpandType, Output = <Self::Output as CubeType>::ExpandType>,
14    >
15{
16    fn __expand_index<'this>(
17        scope: &Scope,
18        this: &'this Self::ExpandType,
19        index: I::ExpandType,
20    ) -> &'this <Self::Output as CubeType>::ExpandType {
21        this.__expand_index_method(scope, index)
22    }
23}
24
25impl<I: CubeType, T: Index<I> + CubeType + ?Sized> CubeIndex<I> for T
26where
27    T::Output: CubeType,
28    T::ExpandType:
29        IndexExpand<I::ExpandType, Output = <<T as Index<I>>::Output as CubeType>::ExpandType>,
30{
31}
32
33pub trait IndexExpand<I> {
34    type Output;
35    fn __expand_index_method(&self, scope: &Scope, index: I) -> &Self::Output;
36}
37
38pub trait CubeIndexMut<I: CubeType>:
39    CubeIndex<I>
40    + IndexMut<I>
41    + CubeType<
42        ExpandType: IndexMutExpand<I::ExpandType, Output = <Self::Output as CubeType>::ExpandType>,
43    >
44{
45    fn __expand_index_mut<'this>(
46        scope: &Scope,
47        this: &'this mut Self::ExpandType,
48        index: I::ExpandType,
49    ) -> &'this mut <Self::Output as CubeType>::ExpandType {
50        this.__expand_index_mut_method(scope, index)
51    }
52}
53
54pub trait IndexMutExpand<I>: IndexExpand<I> {
55    fn __expand_index_mut_method(
56        &mut self,
57        scope: &Scope,
58        index: I,
59    ) -> &mut <Self as IndexExpand<I>>::Output;
60}
61
62impl<I: CubeType, T: IndexMut<I> + CubeIndex<I> + ?Sized> CubeIndexMut<I> for T
63where
64    T::Output: CubeType,
65    T::ExpandType:
66        IndexMutExpand<I::ExpandType, Output = <<T as Index<I>>::Output as CubeType>::ExpandType>,
67{
68}
69
70pub(crate) fn expand_index_native<'a, O>(
71    scope: &Scope,
72    list: Value,
73    index: NativeExpand<usize>,
74    checked: bool,
75) -> &'a O
76where
77    O: From<ExpandValue> + 'static,
78{
79    let index = index.read_value(scope);
80    let val: ExpandValue = index_expand(scope, list, index, checked).into();
81
82    scope.create_kernel_ref(val.into())
83}
84
85pub(crate) fn expand_index_mut_native<'a, O>(
86    scope: &Scope,
87    list: Value,
88    index: NativeExpand<usize>,
89    checked: bool,
90) -> &'a mut O
91where
92    O: From<ExpandValue> + 'static,
93{
94    let index = index.read_value(scope);
95
96    let index_op = IndexOp::maybe_checked(scope.ctx_mut(), list, index, checked);
97    let out: ExpandValue = scope.register_with_result(&index_op).into();
98
99    scope.create_kernel_ref(out.into())
100}
cubecl_core/frontend/indexation.rs

cubecl_core/frontend/
indexation.rs
