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

cubecl_core/frontend/element/
uint.rs
