cubecl_core/frontend/element/
uint.rs

1use cubecl_ir::{ConstantValue, ElemType, Scope, UIntKind, types::scalar::IndexType};
2use pliron::{
3    builtin::types::{IntegerType, Signedness},
4    r#type::TypeHandle,
5};
6
7use crate::prelude::*;
8
9use super::{IntoMut, IntoRuntime, NativeAssign, NativeExpand};
10
11macro_rules! declare_uint {
12    ($primitive:ident, $kind: ident) => {
13        impl CubeType for $primitive {
14            type ExpandType = NativeExpand<Self>;
15        }
16
17        impl Scalar for $primitive {
18            fn elem_type_native() -> ElemType {
19                UIntKind::$kind.into()
20            }
21        }
22        impl CubeDebug for $primitive {}
23        impl CubePrimitive for $primitive {
24            type Scalar = Self;
25            type Size = Const<1>;
26            type WithScalar<S: Scalar> = S;
27
28            fn from_const_value(value: ConstantValue) -> Self {
29                let ConstantValue::UInt(value) = value else {
30                    unreachable!()
31                };
32                value as $primitive
33            }
34        }
35        impl ConstantTyped for $primitive {
36            fn __expand_as_type(scope: &Scope) -> TypeHandle {
37                let width = UIntKind::$kind.size_bits() as u32;
38                IntegerType::get(scope.ctx(), width, Signedness::Unsigned).into()
39            }
40        }
41
42        impl IntoRuntime for $primitive {
43            fn __expand_runtime_method(self, _scope: &Scope) -> NativeExpand<Self> {
44                self.into()
45            }
46        }
47
48        impl IntoExpand for $primitive {
49            type Expand = NativeExpand<$primitive>;
50
51            fn into_expand(self, _: &Scope) -> Self::Expand {
52                self.into()
53            }
54        }
55
56        impl IntoMut for $primitive {
57            fn into_mut(self, _scope: &Scope) -> Self {
58                self
59            }
60        }
61
62        impl NativeAssign for $primitive {}
63
64        impl Numeric for $primitive {
65            fn min_value() -> Self {
66                $primitive::MIN
67            }
68            fn max_value() -> Self {
69                $primitive::MAX
70            }
71        }
72
73        impl Int for $primitive {
74            const BITS: u32 = $primitive::BITS;
75
76            fn new(val: i64) -> Self {
77                val as $primitive
78            }
79        }
80
81        impl_scalar_launch!($primitive);
82    };
83}
84
85declare_uint!(u8, U8);
86declare_uint!(u16, U16);
87declare_uint!(u32, U32);
88declare_uint!(u64, U64);
89
90impl CubeType for usize {
91    type ExpandType = NativeExpand<Self>;
92}
93
94impl CubeDebug for usize {}
95impl Scalar for usize {
96    fn elem_type(_scope: &Scope) -> ElemType {
97        ElemType::Index
98    }
99}
100impl CubePrimitive for usize {
101    type Scalar = Self;
102    type Size = Const<1>;
103    type WithScalar<S: Scalar> = S;
104
105    fn from_const_value(value: ConstantValue) -> Self {
106        let ConstantValue::UInt(value) = value else {
107            unreachable!()
108        };
109        value as usize
110    }
111}
112impl ConstantTyped for usize {
113    fn __expand_as_type(scope: &Scope) -> TypeHandle {
114        IndexType::get(scope.ctx()).into()
115    }
116}
117
118impl IntoRuntime for usize {
119    fn __expand_runtime_method(self, scope: &Scope) -> NativeExpand<Self> {
120        NativeExpand::from_lit(scope, self)
121    }
122}
123
124impl IntoExpand for usize {
125    type Expand = NativeExpand<usize>;
126
127    fn into_expand(self, scope: &Scope) -> Self::Expand {
128        self.__expand_runtime_method(scope)
129    }
130}
131
132impl IntoMut for usize {
133    fn into_mut(self, _scope: &Scope) -> Self {
134        self
135    }
136}
137
138impl NativeAssign for usize {}
139
140impl Numeric for usize {
141    fn min_value() -> Self {
142        usize::MIN
143    }
144    fn max_value() -> Self {
145        // Stay in safe range. Should use runtime version taking scope for correct value.
146        u32::MAX as usize
147    }
148}
149
150impl Int for usize {
151    const BITS: u32 = usize::BITS;
152
153    fn new(val: i64) -> Self {
154        val as usize
155    }
156}
157
158impl_scalar_launch!(usize);
159
160impl CubeType for isize {
161    type ExpandType = NativeExpand<Self>;
162}
163
164impl CubeDebug for isize {}
165impl Scalar for isize {
166    fn elem_type(scope: &Scope) -> ElemType {
167        scope
168            .resolve_type::<Self>()
169            .expect("Type should be registered")
170    }
171}
172impl CubePrimitive for isize {
173    type Scalar = Self;
174    type Size = Const<1>;
175    type WithScalar<S: Scalar> = S;
176
177    fn from_const_value(value: ConstantValue) -> Self {
178        let ConstantValue::Int(value) = value else {
179            unreachable!()
180        };
181        value as isize
182    }
183}
184impl ConstantTyped for isize {
185    fn __expand_as_type(scope: &Scope) -> TypeHandle {
186        scope
187            .resolve_type::<Self>()
188            .expect("Type to be registered")
189            .to_type(scope.ctx_mut())
190    }
191}
192
193impl IntoRuntime for isize {
194    fn __expand_runtime_method(self, scope: &Scope) -> NativeExpand<Self> {
195        NativeExpand::from_lit(scope, self)
196    }
197}
198
199impl IntoExpand for isize {
200    type Expand = NativeExpand<isize>;
201
202    fn into_expand(self, scope: &Scope) -> Self::Expand {
203        self.__expand_runtime_method(scope)
204    }
205}
206
207impl IntoMut for isize {
208    fn into_mut(self, _scope: &Scope) -> Self {
209        self
210    }
211}
212
213impl NativeAssign for isize {}
214
215impl Numeric for isize {
216    fn min_value() -> Self {
217        i32::MIN as isize
218    }
219    fn max_value() -> Self {
220        // Stay in safe range. Should use runtime version taking scope for correct value.
221        i32::MAX as isize
222    }
223}
224
225impl Int for isize {
226    const BITS: u32 = isize::BITS;
227
228    fn new(val: i64) -> Self {
229        val as isize
230    }
231}
232
233impl_scalar_launch!(isize);
cubecl_core/frontend/element/uint.rs

cubecl_core/frontend/element/
uint.rs
