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: &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 CubeDebug for $type {}
52        impl Scalar for $type {}
53        impl CubePrimitive for $type {
54            type Scalar = Self;
55            type Size = Const<1>;
56            type WithScalar<S: Scalar> = S;
57
58            fn as_type_native() -> Option<Type> {
59                Some(ElemType::Int(IntKind::$kind).into())
60            }
61
62            fn from_const_value(value: ConstantValue) -> Self {
63                let ConstantValue::Int(value) = value else {
64                    unreachable!()
65                };
66                value as $type
67            }
68        }
69
70        impl IntoRuntime for $type {
71            fn __expand_runtime_method(self, _scope: &Scope) -> NativeExpand<Self> {
72                self.into()
73            }
74        }
75
76        impl IntoExpand for $type {
77            type Expand = NativeExpand<$type>;
78
79            fn into_expand(self, _: &Scope) -> Self::Expand {
80                self.into()
81            }
82        }
83
84        impl IntoMut for $type {
85            fn into_mut(self, _scope: &Scope) -> Self {
86                self
87            }
88        }
89
90        impl Numeric for $type {
91            fn min_value() -> Self {
92                $type::MIN
93            }
94            fn max_value() -> Self {
95                $type::MAX
96            }
97        }
98
99        impl NativeAssign for $type {}
100
101        impl Int for $type {
102            const BITS: u32 = $type::BITS;
103
104            fn new(val: i64) -> Self {
105                val as $type
106            }
107        }
108
109        impl_scalar_launch!($type);
110    };
111}
112
113impl_int!(i8, I8);
114impl_int!(i16, I16);
115impl_int!(i32, I32);
116impl_int!(i64, I64);
cubecl_core/frontend/element/int.rs

cubecl_core/frontend/element/
int.rs
