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 from_const_value(value: ConstantValue) -> Self {
73                let ConstantValue::Int(value) = value else {
74                    unreachable!()
75                };
76                value as $type
77            }
78        }
79
80        impl ConstantTyped for $type {
81            fn __expand_as_type(scope: &Scope) -> TypeHandle {
82                let width = IntKind::$kind.size_bits() as u32;
83                IntegerType::get(scope.ctx(), width, Signedness::Signed).into()
84            }
85        }
86
87        impl IntoRuntime for $type {
88            fn __expand_runtime_method(self, _scope: &Scope) -> NativeExpand<Self> {
89                self.into()
90            }
91        }
92
93        impl IntoExpand for $type {
94            type Expand = NativeExpand<$type>;
95
96            fn into_expand(self, _: &Scope) -> Self::Expand {
97                self.into()
98            }
99        }
100
101        impl IntoMut for $type {
102            fn into_mut(self, _scope: &Scope) -> Self {
103                self
104            }
105        }
106
107        impl Numeric for $type {
108            fn min_value() -> Self {
109                $type::MIN
110            }
111            fn max_value() -> Self {
112                $type::MAX
113            }
114        }
115
116        impl NativeAssign for $type {}
117
118        impl Int for $type {
119            const BITS: u32 = $type::BITS;
120
121            fn new(val: i64) -> Self {
122                val as $type
123            }
124        }
125
126        impl_scalar_launch!($type);
127    };
128}
129
130impl_int!(i8, I8);
131impl_int!(i16, I16);
132impl_int!(i32, I32);
133impl_int!(i64, I64);
cubecl_core/frontend/element/int.rs

cubecl_core/frontend/element/
int.rs
