cubecl_core/frontend/element/
int.rs

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

cubecl_core/frontend/element/
int.rs
