Skip to main content

cubecl_core/frontend/element/
uint.rs

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