cubecl_core/frontend/element/
cube_elem.rs

1use core::fmt::Debug;
2
3use crate::{
4    self as cubecl, Assign, IntoRuntime,
5    prelude::{Const, CubeDebug, IntoMut, Size},
6};
7use cubecl_ir::{ConstantValue, StorageType, Type, Value, features::TypeUsage};
8use cubecl_macros::{comptime_type, cube, intrinsic};
9use cubecl_runtime::{client::ComputeClient, runtime::Runtime};
10use enumset::EnumSet;
11
12use crate::frontend::CubeType;
13use crate::ir::Scope;
14
15use super::{NativeAssign, NativeExpand};
16
17/// Form of `CubeType` that encapsulates all primitive types:
18/// Numeric, `UInt`, Bool
19pub trait CubePrimitive:
20    CubeType<ExpandType = NativeExpand<Self>>
21    + NativeAssign
22    + CubeDebug
23    + Send
24    + Sync
25    + 'static
26    + Clone
27    + Copy
28{
29    type Scalar: Scalar;
30    type Size: Size;
31    type WithScalar<S: Scalar>: CubePrimitive;
32
33    /// Return the element type to use on GPU.
34    fn as_type() -> Type {
35        Self::as_type_native_unchecked()
36    }
37
38    /// Native or static element type.
39    fn as_type_native() -> Option<Type> {
40        None
41    }
42
43    /// Native or static element type.
44    fn as_type_native_unchecked() -> Type {
45        Self::as_type_native().expect("To be a native type")
46    }
47
48    /// Only native element types have a size.
49    fn size() -> Option<usize> {
50        Self::as_type_native().map(|t| t.size())
51    }
52
53    /// Only native element types have a size.
54    fn size_bits() -> Option<usize> {
55        Self::as_type_native().map(|t| t.size_bits())
56    }
57
58    /// Only native element types have a size.
59    fn size_bits_unchecked() -> usize {
60        Self::as_type_native_unchecked().size_bits()
61    }
62
63    fn from_expand_elem(elem: Value) -> Self::ExpandType {
64        NativeExpand::new(elem)
65    }
66
67    fn from_const_value(value: ConstantValue) -> Self;
68
69    fn into_lit_unchecked(self) -> Self {
70        self
71    }
72
73    fn supported_uses<R: Runtime>(client: &ComputeClient<R>) -> EnumSet<TypeUsage> {
74        let elem = Self::as_type_native_unchecked();
75        client.features().type_usage(elem.storage_type())
76    }
77
78    fn type_size() -> usize {
79        Self::as_type_native_unchecked().size()
80    }
81
82    fn type_size_bits() -> usize {
83        Self::as_type_native_unchecked().size_bits()
84    }
85
86    fn packing_factor() -> usize {
87        Self::as_type_native_unchecked().packing_factor()
88    }
89
90    fn vector_size() -> usize {
91        Self::as_type_native_unchecked().vector_size()
92    }
93
94    fn __expand_as_type(_scope: &Scope) -> Type {
95        Self::as_type_native().expect("To be overridden if not native")
96    }
97
98    fn __expand_type_size(scope: &Scope) -> usize {
99        Self::__expand_as_type(scope).size()
100    }
101
102    fn __expand_type_size_bits(scope: &Scope) -> usize {
103        Self::__expand_as_type(scope).size_bits()
104    }
105
106    fn __expand_packing_factor(scope: &Scope) -> usize {
107        Self::__expand_as_type(scope).packing_factor()
108    }
109
110    fn __expand_vector_size(scope: &Scope) -> usize {
111        Self::__expand_as_type(scope).vector_size()
112    }
113}
114
115pub trait CubePrimitiveExpand {
116    type Scalar: Clone + IntoMut + CubeDebug + Assign;
117    type WithScalar<S: Scalar>: Clone + IntoMut + CubeDebug + Assign;
118}
119
120impl<T: CubePrimitive> CubePrimitiveExpand for NativeExpand<T> {
121    type Scalar = NativeExpand<T::Scalar>;
122    type WithScalar<S: Scalar> = NativeExpand<T::WithScalar<S>>;
123}
124
125/// Marker trait for scalar primitives. Should be implemented for all scalar `CubePrimitive`s, but
126/// **not** for `Vector` or non-standard primitives like `Barrier`. Alternatively, treat these as
127/// types that can be stored in a [`Vector`]
128pub trait Scalar:
129    CubePrimitive<Scalar = Self, Size = Const<1>> + Default + IntoRuntime + Debug + core::cmp::PartialEq
130{
131}
132
133#[cube]
134pub fn type_of<E: CubePrimitive>() -> comptime_type!(Type) {
135    intrinsic!(|scope| { E::__expand_as_type(scope) })
136}
137
138#[cube]
139pub fn storage_type_of<E: CubePrimitive>() -> comptime_type!(StorageType) {
140    intrinsic!(|scope| { E::__expand_as_type(scope).storage_type() })
141}
cubecl_core/frontend/element/cube_elem.rs

cubecl_core/frontend/element/
cube_elem.rs
