cubecl_core/frontend/element/
cube_elem.rs

1use core::fmt::Debug;
2
3use crate::{
4    self as cubecl, Assign, IntoRuntime,
5    frontend::{CanReadValue, IntoExpand, RuntimeAssign, ScalarPartialEq},
6    prelude::{Const, CubeDebug, IntoMut, Size},
7    unexpanded,
8};
9use cubecl_ir::EnumSet;
10use cubecl_ir::{ConstantValue, ElemType, ExpandValue, features::TypeUsage, interfaces::TypedExt};
11use cubecl_macros::{comptime_type, cube, intrinsic};
12use cubecl_runtime::client::Client;
13use pliron::r#type::TypeHandle;
14
15use crate::frontend::CubeType;
16use crate::ir::Scope;
17
18use super::{NativeAssign, NativeExpand};
19
20/// Form of `CubeType` that encapsulates all primitive types:
21/// Numeric, `UInt`, Bool
22pub trait CubePrimitive:
23    CubeType<ExpandType = NativeExpand<Self>>
24    + NativeAssign
25    + ConstantTyped
26    + CanReadValue
27    + CubeDebug
28    + Send
29    + Sync
30    + 'static
31    + Clone
32    + Copy
33{
34    type Scalar: Scalar;
35    type Size: Size;
36    type WithScalar<S: Scalar>: CubePrimitive;
37
38    fn from_expand_elem(elem: ExpandValue) -> Self::ExpandType {
39        NativeExpand::new(elem)
40    }
41
42    fn from_const_value(value: ConstantValue) -> Self;
43
44    fn into_lit_unchecked(self) -> Self {
45        self
46    }
47
48    fn packing_factor() -> usize {
49        unexpanded!()
50    }
51
52    fn vector_size() -> usize {
53        unexpanded!()
54    }
55
56    fn __expand_packing_factor(scope: &Scope) -> usize {
57        Self::__expand_as_type(scope).packing_factor(scope.ctx())
58    }
59
60    fn __expand_vector_size(scope: &Scope) -> usize {
61        Self::__expand_as_type(scope).vector_size(scope.ctx())
62    }
63}
64
65/// Types that contain all the information needed to construct a runtime type. Implemented for most
66/// native types, but not `Array<T>` because it has a size defined at construction.
67pub trait ConstantTyped: CubeType<ExpandType = NativeExpand<Self>> + Sized {
68    /// Return the element type to use on GPU.
69    fn as_type() -> TypeHandle {
70        unexpanded!()
71    }
72
73    /// Only native element types have a size.
74    fn size() -> usize {
75        size_of::<Self>()
76    }
77
78    /// Only native element types have a size.
79    fn size_bits() -> usize {
80        Self::size() * 8
81    }
82
83    fn __expand_as_type(scope: &Scope) -> TypeHandle;
84
85    fn __expand_size(scope: &Scope) -> usize {
86        Self::__expand_as_type(scope).size(scope.ctx())
87    }
88
89    fn __expand_size_bits(scope: &Scope) -> usize {
90        Self::__expand_size(scope) * 8
91    }
92}
93
94/// `RuntimeAssign` for `ConstantTyped` types, doesn't need a reference element.
95pub trait ConstRuntimeAssign<T = <Self as IntoExpand>::Expand>: RuntimeAssign<T> {
96    /// Create a new mutable variable of this type in `scope`.
97    fn const_init_mut(scope: &Scope) -> Self::Expand;
98}
99
100pub trait CubePrimitiveExpand {
101    type Scalar: Clone + IntoMut + CubeDebug + Assign;
102    type WithScalar<S: Scalar>: Clone + IntoMut + CubeDebug + Assign;
103}
104
105impl<T: CubePrimitive> CubePrimitiveExpand for NativeExpand<T> {
106    type Scalar = NativeExpand<T::Scalar>;
107    type WithScalar<S: Scalar> = NativeExpand<T::WithScalar<S>>;
108}
109
110/// Marker trait for scalar primitives. Should be implemented for all scalar `CubePrimitive`s, but
111/// **not** for `Vector` or non-standard primitives like `Barrier`. Alternatively, treat these as
112/// types that can be stored in a [`Vector`](crate::prelude::Vector)
113pub trait Scalar:
114    CubePrimitive<Scalar = Self, Size = Const<1>, WithScalar<Self> = Self>
115    + Default
116    + IntoRuntime
117    + Debug
118    + ScalarPartialEq
119    + Into<ExpandValue>
120{
121    fn elem_type(_scope: &Scope) -> ElemType {
122        Self::elem_type_native()
123    }
124    fn elem_type_native() -> ElemType {
125        unexpanded!()
126    }
127
128    fn supported_uses(client: &Client) -> EnumSet<TypeUsage> {
129        let ty = Self::elem_type_native();
130        client.features().type_usage(ty)
131    }
132}
133
134impl CubeDebug for TypeHandle {}
135impl CubeDebug for ElemType {}
136
137#[cube]
138pub fn elem_type_of<E: CubePrimitive>() -> comptime_type!(ElemType) {
139    intrinsic!(|scope| { E::Scalar::elem_type(scope) })
140}
cubecl_core/frontend/element/cube_elem.rs

cubecl_core/frontend/element/
cube_elem.rs
