cubecl_core/frontend/
indexation.rs

1use core::ops::{Index, IndexMut};
2
3use cubecl_ir::{IndexOperands, Instruction, Memory, Scope, Type, Value, ValueKind};
4
5use super::{CubeType, NativeExpand, index_expand};
6use crate::prelude::CubePrimitive;
7
8/// Trait bound that can be used to guarantee the expand also implements `IndexExpand`
9pub trait CubeIndex<I: CubeType>:
10    Index<I, Output: CubeType>
11    + CubeType<
12        ExpandType: IndexExpand<I::ExpandType, Output = <Self::Output as CubeType>::ExpandType>,
13    >
14{
15    fn __expand_index<'this>(
16        scope: &Scope,
17        this: &'this Self::ExpandType,
18        index: I::ExpandType,
19    ) -> &'this <Self::Output as CubeType>::ExpandType {
20        this.__expand_index_method(scope, index)
21    }
22}
23
24impl<I: CubeType, T: Index<I> + CubeType + ?Sized> CubeIndex<I> for T
25where
26    T::Output: CubeType,
27    T::ExpandType:
28        IndexExpand<I::ExpandType, Output = <<T as Index<I>>::Output as CubeType>::ExpandType>,
29{
30}
31
32pub trait IndexExpand<I> {
33    type Output;
34    fn __expand_index_method(&self, scope: &Scope, index: I) -> &Self::Output;
35}
36
37pub trait CubeIndexMut<I: CubeType>:
38    CubeIndex<I>
39    + IndexMut<I>
40    + CubeType<
41        ExpandType: IndexMutExpand<I::ExpandType, Output = <Self::Output as CubeType>::ExpandType>,
42    >
43{
44    fn __expand_index_mut<'this>(
45        scope: &Scope,
46        this: &'this mut Self::ExpandType,
47        index: I::ExpandType,
48    ) -> &'this mut <Self::Output as CubeType>::ExpandType {
49        this.__expand_index_mut_method(scope, index)
50    }
51}
52
53pub trait IndexMutExpand<I>: IndexExpand<I> {
54    fn __expand_index_mut_method(
55        &mut self,
56        scope: &Scope,
57        index: I,
58    ) -> &mut <Self as IndexExpand<I>>::Output;
59}
60
61impl<I: CubeType, T: IndexMut<I> + CubeIndex<I> + ?Sized> CubeIndexMut<I> for T
62where
63    T::Output: CubeType,
64    T::ExpandType:
65        IndexMutExpand<I::ExpandType, Output = <<T as Index<I>>::Output as CubeType>::ExpandType>,
66{
67}
68
69pub(crate) fn expand_index_native<'a, O>(
70    scope: &Scope,
71    list: Value,
72    index: NativeExpand<usize>,
73    checked: bool,
74) -> &'a O
75where
76    O: From<Value> + 'static,
77{
78    let index: Value = index.into();
79    let index_var: Value = index;
80    let index = match index_var.kind {
81        ValueKind::Constant(value) => Value::constant(value, usize::__expand_as_type(scope)),
82        _ => index,
83    };
84    let val = index_expand(scope, list, index, checked);
85
86    scope.create_kernel_ref(val.into())
87}
88
89pub(crate) fn expand_index_mut_native<'a, O>(
90    scope: &Scope,
91    list: Value,
92    index: NativeExpand<usize>,
93    checked: bool,
94) -> &'a mut O
95where
96    O: From<Value> + 'static,
97{
98    let index: Value = index.expand;
99    let index = match index.kind {
100        ValueKind::Constant(value) => Value::constant(value, usize::__expand_as_type(scope)),
101        _ => index,
102    };
103
104    let ty = list.value_type();
105    let class = list.address_space();
106    let out = scope.create_value(Type::pointer(ty, class));
107
108    scope.register(Instruction::new(
109        Memory::Index(IndexOperands {
110            list,
111            index,
112            unroll_factor: 1,
113            checked,
114        }),
115        out,
116    ));
117
118    scope.create_kernel_ref(out.into())
119}
cubecl_core/frontend/indexation.rs

cubecl_core/frontend/
indexation.rs
