cubecl_core/frontend/
indexation.rs

1use cubecl_ir::ExpandElement;
2
3use super::{CubeType, ExpandElementTyped};
4use crate::{
5    ir::{IntKind, UIntKind, Variable},
6    unexpanded,
7};
8
9/// Fake indexation so we can rewrite indexes into scalars as calls to this fake function in the
10/// non-expanded function
11pub trait CubeIndex<T: Index> {
12    type Output: CubeType;
13
14    fn cube_idx(&self, _i: T) -> &Self::Output {
15        unexpanded!()
16    }
17}
18
19pub trait CubeIndexMut<T: Index>: CubeIndex<T> {
20    fn cube_idx_mut(&mut self, _i: T) -> &mut Self::Output {
21        unexpanded!()
22    }
23}
24
25pub trait Index {
26    fn value(self) -> Variable;
27}
28
29impl Index for i32 {
30    fn value(self) -> Variable {
31        Variable::constant(crate::ir::ConstantScalarValue::Int(
32            self as i64,
33            IntKind::I32,
34        ))
35    }
36}
37
38impl Index for u32 {
39    fn value(self) -> Variable {
40        Variable::constant(crate::ir::ConstantScalarValue::UInt(
41            self as u64,
42            UIntKind::U32,
43        ))
44    }
45}
46
47impl Index for ExpandElement {
48    fn value(self) -> Variable {
49        *self
50    }
51}
52
53impl Index for ExpandElementTyped<u32> {
54    fn value(self) -> Variable {
55        *self.expand
56    }
57}
cubecl_core/frontend/indexation.rs

cubecl_core/frontend/
indexation.rs
