cubecl_core/frontend/element/float/
fp8.rs

1use cubecl_common::{e4m3, e5m2, ue8m0};
2use cubecl_ir::{ConstantValue, ElemType, ExpandElement, FloatKind, Scope, StorageType};
3
4use crate::prelude::{
5    CubePrimitive, CubeType, ExpandElementIntoMut, ExpandElementTyped, IntoRuntime,
6    into_mut_expand_element, into_runtime_expand_element,
7};
8
9impl CubeType for e4m3 {
10    type ExpandType = ExpandElementTyped<e4m3>;
11}
12
13impl CubePrimitive for e4m3 {
14    /// Return the element type to use on GPU
15    fn as_type_native() -> Option<StorageType> {
16        Some(ElemType::Float(FloatKind::E4M3).into())
17    }
18
19    fn from_const_value(value: ConstantValue) -> Self {
20        let ConstantValue::Float(value) = value else {
21            unreachable!()
22        };
23        e4m3::from_f64(value)
24    }
25}
26
27impl IntoRuntime for e4m3 {
28    fn __expand_runtime_method(self, scope: &mut Scope) -> ExpandElementTyped<Self> {
29        let elem: ExpandElementTyped<Self> = self.into();
30        into_runtime_expand_element(scope, elem).into()
31    }
32}
33
34impl ExpandElementIntoMut for e4m3 {
35    fn elem_into_mut(scope: &mut Scope, elem: ExpandElement) -> ExpandElement {
36        into_mut_expand_element(scope, elem)
37    }
38}
39
40impl CubeType for e5m2 {
41    type ExpandType = ExpandElementTyped<e5m2>;
42}
43
44impl CubePrimitive for e5m2 {
45    /// Return the element type to use on GPU
46    fn as_type_native() -> Option<StorageType> {
47        Some(ElemType::Float(FloatKind::E5M2).into())
48    }
49
50    fn from_const_value(value: ConstantValue) -> Self {
51        let ConstantValue::Float(value) = value else {
52            unreachable!()
53        };
54        e5m2::from_f64(value)
55    }
56}
57
58impl IntoRuntime for e5m2 {
59    fn __expand_runtime_method(self, scope: &mut Scope) -> ExpandElementTyped<Self> {
60        let elem: ExpandElementTyped<Self> = self.into();
61        into_runtime_expand_element(scope, elem).into()
62    }
63}
64
65impl ExpandElementIntoMut for e5m2 {
66    fn elem_into_mut(scope: &mut Scope, elem: ExpandElement) -> ExpandElement {
67        into_mut_expand_element(scope, elem)
68    }
69}
70
71impl CubeType for ue8m0 {
72    type ExpandType = ExpandElementTyped<ue8m0>;
73}
74
75impl CubePrimitive for ue8m0 {
76    /// Return the element type to use on GPU
77    fn as_type_native() -> Option<StorageType> {
78        Some(ElemType::Float(FloatKind::UE8M0).into())
79    }
80
81    fn from_const_value(value: ConstantValue) -> Self {
82        let ConstantValue::Float(value) = value else {
83            unreachable!()
84        };
85        ue8m0::from_f64(value)
86    }
87}
88
89impl IntoRuntime for ue8m0 {
90    fn __expand_runtime_method(self, scope: &mut Scope) -> ExpandElementTyped<Self> {
91        let elem: ExpandElementTyped<Self> = self.into();
92        into_runtime_expand_element(scope, elem).into()
93    }
94}
95
96impl ExpandElementIntoMut for ue8m0 {
97    fn elem_into_mut(scope: &mut Scope, elem: ExpandElement) -> ExpandElement {
98        into_mut_expand_element(scope, elem)
99    }
100}
cubecl_core/frontend/element/float/fp8.rs

cubecl_core/frontend/element/float/
fp8.rs
