cubecl_core/frontend/element/
cube_elem.rs

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

cubecl_core/frontend/element/
cube_elem.rs
