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, TrailingZeros};
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    + TrailingZeros
27    + FindFirstSet
28    + SaturatingAdd
29    + SaturatingSub
30    + core::ops::BitOr<Output = Self>
31    + core::ops::BitAnd<Output = Self>
32    + core::ops::BitXor<Output = Self>
33    + core::ops::Shl<Output = Self>
34    + core::ops::Shr<Output = Self>
35    + core::ops::Not<Output = Self>
36    + core::ops::BitOrAssign
37    + core::ops::BitAndAssign
38    + core::ops::BitXorAssign
39    + core::ops::ShlAssign<u32>
40    + core::ops::ShrAssign<u32>
41    + core::hash::Hash
42    + core::cmp::PartialOrd
43    + core::cmp::Ord
44    + core::cmp::PartialEq
45    + core::cmp::Eq
46{
47    const BITS: u32;
48
49    fn new(val: i64) -> Self;
50    fn __expand_new(scope: &mut Scope, val: i64) -> <Self as CubeType>::ExpandType {
51        __expand_new(scope, val)
52    }
53}
54
55macro_rules! impl_int {
56    ($type:ident, $kind:ident) => {
57        impl CubeType for $type {
58            type ExpandType = ExpandElementTyped<Self>;
59        }
60
61        impl CubePrimitive for $type {
62            fn as_type_native() -> Option<StorageType> {
63                Some(ElemType::Int(IntKind::$kind).into())
64            }
65
66            fn from_const_value(value: ConstantValue) -> Self {
67                let ConstantValue::Int(value) = value else {
68                    unreachable!()
69                };
70                value as $type
71            }
72        }
73
74        impl IntoRuntime for $type {
75            fn __expand_runtime_method(self, scope: &mut Scope) -> ExpandElementTyped<Self> {
76                let elem: ExpandElementTyped<Self> = self.into();
77                into_runtime_expand_element(scope, elem).into()
78            }
79        }
80
81        impl IntoMut for $type {
82            fn into_mut(self, _scope: &mut Scope) -> Self {
83                self
84            }
85        }
86
87        impl Numeric for $type {
88            fn min_value() -> Self {
89                $type::MIN
90            }
91            fn max_value() -> Self {
92                $type::MAX
93            }
94        }
95
96        impl ExpandElementIntoMut for $type {
97            fn elem_into_mut(scope: &mut Scope, elem: ExpandElement) -> ExpandElement {
98                into_mut_expand_element(scope, elem)
99            }
100        }
101
102        impl Int for $type {
103            const BITS: u32 = $type::BITS;
104
105            fn new(val: i64) -> Self {
106                val as $type
107            }
108        }
109    };
110}
111
112impl_int!(i8, I8);
113impl_int!(i16, I16);
114impl_int!(i32, I32);
115impl_int!(i64, I64);
cubecl_core/frontend/element/int.rs

cubecl_core/frontend/element/
int.rs
