cubecl_core/frontend/element/
int.rs

1use cubecl_ir::{ConstantValue, Type};
2
3use crate::frontend::{CubeType, Numeric};
4use crate::ir::{ElemType, IntKind, Scope};
5use crate::prelude::*;
6
7use super::{__expand_new, CubePrimitive, IntoMut, IntoRuntime, NativeAssign, NativeExpand};
8
9/// Signed or unsigned integer. Used as input in int kernels
10pub trait Int:
11    Numeric
12    + CubeNot
13    + CountOnes
14    + ReverseBits
15    + LeadingZeros
16    + TrailingZeros
17    + FindFirstSet
18    + SaturatingAdd
19    + SaturatingSub
20    + core::ops::BitOr<Output = Self>
21    + core::ops::BitAnd<Output = Self>
22    + core::ops::BitXor<Output = Self>
23    + core::ops::Shl<Output = Self>
24    + core::ops::Shr<Output = Self>
25    + core::ops::Not<Output = Self>
26    + core::ops::BitOrAssign
27    + core::ops::BitAndAssign
28    + core::ops::BitXorAssign
29    + core::ops::ShlAssign<u32>
30    + core::ops::ShrAssign<u32>
31    + core::hash::Hash
32    + core::cmp::PartialOrd
33    + core::cmp::Ord
34    + core::cmp::PartialEq
35    + core::cmp::Eq
36{
37    const BITS: u32;
38
39    fn new(val: i64) -> Self;
40    fn __expand_new(scope: &mut Scope, val: i64) -> <Self as CubeType>::ExpandType {
41        __expand_new(scope, val)
42    }
43}
44
45macro_rules! impl_int {
46    ($type:ident, $kind:ident) => {
47        impl CubeType for $type {
48            type ExpandType = NativeExpand<Self>;
49        }
50
51        impl Scalar for $type {}
52        impl CubePrimitive for $type {
53            type Scalar = Self;
54            type Size = Const<1>;
55            type WithScalar<S: Scalar> = S;
56
57            fn as_type_native() -> Option<Type> {
58                Some(ElemType::Int(IntKind::$kind).into())
59            }
60
61            fn from_const_value(value: ConstantValue) -> Self {
62                let ConstantValue::Int(value) = value else {
63                    unreachable!()
64                };
65                value as $type
66            }
67        }
68
69        impl IntoRuntime for $type {
70            fn __expand_runtime_method(self, _scope: &mut Scope) -> NativeExpand<Self> {
71                self.into()
72            }
73        }
74
75        impl IntoMut for $type {
76            fn into_mut(self, _scope: &mut Scope) -> Self {
77                self
78            }
79        }
80
81        impl Numeric for $type {
82            fn min_value() -> Self {
83                $type::MIN
84            }
85            fn max_value() -> Self {
86                $type::MAX
87            }
88        }
89
90        impl NativeAssign for $type {}
91
92        impl Int for $type {
93            const BITS: u32 = $type::BITS;
94
95            fn new(val: i64) -> Self {
96                val as $type
97            }
98        }
99    };
100}
101
102impl_int!(i8, I8);
103impl_int!(i16, I16);
104impl_int!(i32, I32);
105impl_int!(i64, I64);
cubecl_core/frontend/element/int.rs

cubecl_core/frontend/element/
int.rs
