cubecl_core/frontend/element/
uint.rs

1use cubecl_ir::{ConstantValue, ExpandElement, Scope, StorageType, UIntKind};
2
3use crate::frontend::{CubePrimitive, CubeType, Numeric};
4use crate::ir::ElemType;
5
6use super::{
7    ExpandElementIntoMut, ExpandElementTyped, Int, IntoMut, IntoRuntime, into_mut_expand_element,
8    into_runtime_expand_element,
9};
10
11macro_rules! declare_uint {
12    ($primitive:ident, $kind:ident) => {
13        impl CubeType for $primitive {
14            type ExpandType = ExpandElementTyped<Self>;
15        }
16
17        impl CubePrimitive for $primitive {
18            fn as_type_native() -> Option<StorageType> {
19                Some(ElemType::UInt(UIntKind::$kind).into())
20            }
21
22            fn from_const_value(value: ConstantValue) -> Self {
23                let ConstantValue::UInt(value) = value else {
24                    unreachable!()
25                };
26                value as $primitive
27            }
28        }
29
30        impl IntoRuntime for $primitive {
31            fn __expand_runtime_method(self, scope: &mut Scope) -> ExpandElementTyped<Self> {
32                let elem: ExpandElementTyped<Self> = self.into();
33                into_runtime_expand_element(scope, elem).into()
34            }
35        }
36
37        impl IntoMut for $primitive {
38            fn into_mut(self, _scope: &mut Scope) -> Self {
39                self
40            }
41        }
42
43        impl ExpandElementIntoMut for $primitive {
44            fn elem_into_mut(scope: &mut Scope, elem: ExpandElement) -> ExpandElement {
45                into_mut_expand_element(scope, elem)
46            }
47        }
48
49        impl Numeric for $primitive {
50            fn min_value() -> Self {
51                $primitive::MIN
52            }
53            fn max_value() -> Self {
54                $primitive::MAX
55            }
56        }
57
58        impl Int for $primitive {
59            const BITS: u32 = $primitive::BITS;
60
61            fn new(val: i64) -> Self {
62                val as $primitive
63            }
64        }
65    };
66}
67
68declare_uint!(u8, U8);
69declare_uint!(u16, U16);
70declare_uint!(u32, U32);
71declare_uint!(u64, U64);
72
73impl CubeType for usize {
74    type ExpandType = ExpandElementTyped<Self>;
75}
76
77impl CubePrimitive for usize {
78    fn from_const_value(value: ConstantValue) -> Self {
79        let ConstantValue::UInt(value) = value else {
80            unreachable!()
81        };
82        value as usize
83    }
84
85    fn as_type(scope: &Scope) -> StorageType {
86        scope.resolve_type::<Self>().expect("Type to be registered")
87    }
88}
89
90impl IntoRuntime for usize {
91    fn __expand_runtime_method(self, scope: &mut Scope) -> ExpandElementTyped<Self> {
92        let elem: ExpandElementTyped<Self> = self.into();
93        into_runtime_expand_element(scope, elem).into()
94    }
95}
96
97impl IntoMut for usize {
98    fn into_mut(self, _scope: &mut Scope) -> Self {
99        self
100    }
101}
102
103impl ExpandElementIntoMut for usize {
104    fn elem_into_mut(scope: &mut Scope, elem: ExpandElement) -> ExpandElement {
105        into_mut_expand_element(scope, elem)
106    }
107}
108
109impl Numeric for usize {
110    fn min_value() -> Self {
111        usize::MIN
112    }
113    fn max_value() -> Self {
114        // Stay in safe range. Should use runtime version taking scope for correct value.
115        u32::MAX as usize
116    }
117}
118
119impl Int for usize {
120    const BITS: u32 = usize::BITS;
121
122    fn new(val: i64) -> Self {
123        val as usize
124    }
125}
126
127impl CubeType for isize {
128    type ExpandType = ExpandElementTyped<Self>;
129}
130
131impl CubePrimitive for isize {
132    fn from_const_value(value: ConstantValue) -> Self {
133        let ConstantValue::Int(value) = value else {
134            unreachable!()
135        };
136        value as isize
137    }
138
139    fn as_type(scope: &Scope) -> StorageType {
140        scope.resolve_type::<Self>().expect("Type to be registered")
141    }
142}
143
144impl IntoRuntime for isize {
145    fn __expand_runtime_method(self, scope: &mut Scope) -> ExpandElementTyped<Self> {
146        let elem: ExpandElementTyped<Self> = self.into();
147        into_runtime_expand_element(scope, elem).into()
148    }
149}
150
151impl IntoMut for isize {
152    fn into_mut(self, _scope: &mut Scope) -> Self {
153        self
154    }
155}
156
157impl ExpandElementIntoMut for isize {
158    fn elem_into_mut(scope: &mut Scope, elem: ExpandElement) -> ExpandElement {
159        into_mut_expand_element(scope, elem)
160    }
161}
162
163impl Numeric for isize {
164    fn min_value() -> Self {
165        i32::MIN as isize
166    }
167    fn max_value() -> Self {
168        // Stay in safe range. Should use runtime version taking scope for correct value.
169        i32::MAX as isize
170    }
171}
172
173impl Int for isize {
174    const BITS: u32 = isize::BITS;
175
176    fn new(val: i64) -> Self {
177        val as isize
178    }
179}
cubecl_core/frontend/element/uint.rs

cubecl_core/frontend/element/
uint.rs
