cubecl_core/frontend/
indexation.rs

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

cubecl_core/frontend/
indexation.rs
