cubecl_core/frontend/element/
cube_elem.rs

1use crate as cubecl;
2use cubecl_ir::{ConstantScalarValue, ExpandElement, StorageType};
3use cubecl_macros::{comptime_type, cube, intrinsic};
4use cubecl_runtime::{TypeUsage, client::ComputeClient, runtime::Runtime};
5use enumset::EnumSet;
6
7use crate::frontend::CubeType;
8use crate::ir::Scope;
9
10use super::{ExpandElementIntoMut, ExpandElementTyped};
11
12/// Form of CubeType that encapsulates all primitive types:
13/// Numeric, UInt, Bool
14pub trait CubePrimitive:
15    CubeType<ExpandType = ExpandElementTyped<Self>>
16    + ExpandElementIntoMut
17    // + IntoRuntime
18    + core::cmp::PartialEq
19    + Send
20    + Sync
21    + 'static
22    + Clone
23    + Copy
24{
25    /// Return the element type to use on GPU.
26    fn as_type(_scope: &Scope) -> StorageType {
27        Self::as_type_native().expect("To be overridden if not native")
28    }
29
30    /// Native or static element type.
31    fn as_type_native() -> Option<StorageType> {
32        None
33    }
34
35    /// Native or static element type.
36    fn as_type_native_unchecked() -> StorageType {
37        Self::as_type_native().expect("To be a native type")
38    }
39
40    /// Only native element types have a size.
41    fn size() -> Option<usize> {
42        Self::as_type_native().map(|t| t.size())
43    }
44
45    /// Only native element types have a size.
46    fn size_bits() -> Option<usize> {
47        Self::as_type_native().map(|t| t.size_bits())
48    }
49
50    /// Only native element types have a size.
51    fn size_bits_unchecked() -> usize {
52        Self::as_type_native_unchecked().size_bits()
53    }
54
55    fn from_expand_elem(elem: ExpandElement) -> Self::ExpandType {
56        ExpandElementTyped::new(elem)
57    }
58
59    fn from_const_value(value: ConstantScalarValue) -> Self;
60
61    fn into_lit_unchecked(self) -> Self {
62        self
63    }
64
65    fn supported_uses<R: Runtime>(
66        client: &ComputeClient<R>,
67    ) -> EnumSet<TypeUsage> {
68        let elem = Self::as_type_native_unchecked();
69        client.properties().features.type_usage(elem)
70    }
71
72    fn type_size() -> u32 {
73        Self::as_type_native_unchecked().size() as u32
74    }
75
76    fn type_size_bits() -> u32 {
77        Self::as_type_native_unchecked().size_bits() as u32
78    }
79
80    fn packing_factor() -> u32 {
81        Self::as_type_native_unchecked().packing_factor()
82    }
83
84    fn __expand_type_size(scope: &Scope) -> u32 {
85        Self::as_type(scope).size() as u32
86    }
87
88    fn __expand_type_size_bits(scope: &Scope) -> u32 {
89        Self::as_type(scope).size_bits() as u32
90    }
91
92    fn __expand_packing_factor(scope: &Scope) -> u32 {
93        Self::as_type(scope).packing_factor()
94    }
95}
96
97#[cube]
98pub fn type_of<E: CubePrimitive>() -> comptime_type!(StorageType) {
99    intrinsic!(|scope| { E::as_type(scope) })
100}
cubecl_core/frontend/element/cube_elem.rs

cubecl_core/frontend/element/
cube_elem.rs
