cubecl_core/frontend/element/
uint.rs

1use cubecl_ir::{ExpandElement, Scope, UIntKind};
2
3use crate::Runtime;
4use crate::frontend::{CubePrimitive, CubeType, Numeric};
5use crate::ir::Elem;
6use crate::prelude::{KernelBuilder, KernelLauncher};
7
8use super::{
9    ExpandElementBaseInit, ExpandElementTyped, Init, Int, IntoRuntime, LaunchArgExpand,
10    ScalarArgSettings, init_expand_element,
11};
12
13macro_rules! declare_uint {
14    ($primitive:ident, $kind:ident) => {
15        impl CubeType for $primitive {
16            type ExpandType = ExpandElementTyped<Self>;
17        }
18
19        impl ExpandElementBaseInit for $primitive {
20            fn init_elem(scope: &mut Scope, elem: ExpandElement) -> ExpandElement {
21                init_expand_element(scope, elem)
22            }
23        }
24
25        impl CubePrimitive for $primitive {
26            fn as_elem_native() -> Option<Elem> {
27                Some(Elem::UInt(UIntKind::$kind))
28            }
29        }
30
31        impl Init for $primitive {
32            fn init(self, _scope: &mut Scope) -> Self {
33                self
34            }
35        }
36
37        impl IntoRuntime for $primitive {
38            fn __expand_runtime_method(self, scope: &mut Scope) -> ExpandElementTyped<Self> {
39                let expand: ExpandElementTyped<Self> = self.into();
40                Init::init(expand, scope)
41            }
42        }
43
44        impl LaunchArgExpand for $primitive {
45            type CompilationArg = ();
46
47            fn expand(
48                _: &Self::CompilationArg,
49                builder: &mut KernelBuilder,
50            ) -> ExpandElementTyped<Self> {
51                builder.scalar($primitive::as_elem(&builder.context)).into()
52            }
53        }
54
55        impl Numeric for $primitive {
56            fn min_value() -> Self {
57                $primitive::MIN
58            }
59            fn max_value() -> Self {
60                $primitive::MAX
61            }
62        }
63
64        impl Int for $primitive {
65            const BITS: u32 = $primitive::BITS;
66
67            fn new(val: i64) -> Self {
68                val as $primitive
69            }
70        }
71    };
72}
73
74declare_uint!(u8, U8);
75declare_uint!(u16, U16);
76declare_uint!(u32, U32);
77declare_uint!(u64, U64);
78
79impl ScalarArgSettings for u8 {
80    fn register<R: Runtime>(&self, settings: &mut KernelLauncher<R>) {
81        settings.register_u8(*self);
82    }
83}
84
85impl ScalarArgSettings for u16 {
86    fn register<R: Runtime>(&self, settings: &mut KernelLauncher<R>) {
87        settings.register_u16(*self);
88    }
89}
90
91impl ScalarArgSettings for u32 {
92    fn register<R: Runtime>(&self, settings: &mut KernelLauncher<R>) {
93        settings.register_u32(*self);
94    }
95}
96
97impl ScalarArgSettings for u64 {
98    fn register<R: Runtime>(&self, settings: &mut KernelLauncher<R>) {
99        settings.register_u64(*self);
100    }
101}
cubecl_core/frontend/element/uint.rs

cubecl_core/frontend/element/
uint.rs
