cubecl_core/frontend/element/
int.rs

1use cubecl_ir::{ConstantValue, ElemType};
2use pliron::{
3    builtin::types::{IntegerType, Signedness},
4    r#type::TypeHandle,
5};
6
7use crate::frontend::{CubeType, Numeric};
8use crate::ir::{IntKind, Scope};
9use crate::prelude::*;
10
11use super::{__expand_new, CubePrimitive, IntoMut, IntoRuntime, NativeAssign, NativeExpand};
12
13/// Signed or unsigned integer. Used as input in int kernels
14pub trait Int:
15    Numeric
16    + CubeNot
17    + CountOnes
18    + ScalarReverseBits
19    + LeadingZeros
20    + TrailingZeros
21    + FindFirstSet
22    + ScalarSaturatingAdd
23    + ScalarSaturatingSub
24    + ScalarBitOr
25    + ScalarBitAnd
26    + ScalarBitXor
27    + ScalarShl
28    + ScalarShr
29    + CubeBitOrAssign
30    + CubeBitAndAssign
31    + CubeBitXorAssign
32    + CubeShlAssign
33    + CubeShrAssign
34    + ScalarOrd
35    + core::ops::ShlAssign<u32>
36    + core::ops::ShrAssign<u32>
37    + core::hash::Hash
38    + core::cmp::PartialOrd
39    + core::cmp::Ord
40    + core::cmp::PartialEq
41    + core::cmp::Eq
42{
43    const BITS: u32;
44
45    fn new(val: i64) -> Self;
46    fn __expand_new(scope: &Scope, val: i64) -> <Self as CubeType>::ExpandType {
47        __expand_new(scope, val)
48    }
49
50    fn is_signed(scope: &Scope) -> bool {
51        Self::elem_type(scope).is_signed_int()
52    }
53}
54
55macro_rules! impl_int {
56    ($type: ident, $kind: ident) => {
57        impl CubeType for $type {
58            type ExpandType = NativeExpand<Self>;
59        }
60
61        impl CubeDebug for $type {}
62        impl Scalar for $type {
63            fn elem_type_native() -> ElemType {
64                IntKind::$kind.into()
65            }
66        }
67        impl CubePrimitive for $type {
68            type Scalar = Self;
69            type Size = Const<1>;
70            type WithScalar<S: Scalar> = S;
71
72            fn __expand_as_type(scope: &Scope) -> TypeHandle {
73                let width = IntKind::$kind.size_bits() as u32;
74                IntegerType::get(scope.ctx(), width, Signedness::Signed).into()
75            }
76
77            fn from_const_value(value: ConstantValue) -> Self {
78                let ConstantValue::Int(value) = value else {
79                    unreachable!()
80                };
81                value as $type
82            }
83        }
84
85        impl IntoRuntime for $type {
86            fn __expand_runtime_method(self, _scope: &Scope) -> NativeExpand<Self> {
87                self.into()
88            }
89        }
90
91        impl IntoExpand for $type {
92            type Expand = NativeExpand<$type>;
93
94            fn into_expand(self, _: &Scope) -> Self::Expand {
95                self.into()
96            }
97        }
98
99        impl IntoMut for $type {
100            fn into_mut(self, _scope: &Scope) -> Self {
101                self
102            }
103        }
104
105        impl Numeric for $type {
106            fn min_value() -> Self {
107                $type::MIN
108            }
109            fn max_value() -> Self {
110                $type::MAX
111            }
112        }
113
114        impl NativeAssign for $type {}
115
116        impl Int for $type {
117            const BITS: u32 = $type::BITS;
118
119            fn new(val: i64) -> Self {
120                val as $type
121            }
122        }
123
124        impl_scalar_launch!($type);
125    };
126}
127
128impl_int!(i8, I8);
129impl_int!(i16, I16);
130impl_int!(i32, I32);
131impl_int!(i64, I64);
cubecl_core/frontend/element/int.rs

cubecl_core/frontend/element/
int.rs
