cubecl_core/frontend/element/
int.rs

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

cubecl_core/frontend/element/
int.rs
