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