Skip to main content

cubecl_core/frontend/element/float/
fp8.rs

1use cubecl_common::{e4m3, e5m2, ue8m0};
2use cubecl_ir::{ConstantValue, ElemType, FloatKind, Scope, Type};
3
4use crate::prelude::*;
5
6impl CubeType for e4m3 {
7    type ExpandType = NativeExpand<e4m3>;
8}
9
10impl CubeDebug for e4m3 {}
11impl Scalar for e4m3 {}
12impl CubePrimitive for e4m3 {
13    type Scalar = Self;
14    type Size = Const<1>;
15    type WithScalar<S: Scalar> = S;
16
17    /// Return the element type to use on GPU
18    fn as_type_native() -> Option<Type> {
19        Some(ElemType::Float(FloatKind::E4M3).into())
20    }
21
22    fn from_const_value(value: ConstantValue) -> Self {
23        let ConstantValue::Float(value) = value else {
24            unreachable!()
25        };
26        e4m3::from_f64(value)
27    }
28}
29
30impl IntoRuntime for e4m3 {
31    fn __expand_runtime_method(self, _scope: &Scope) -> NativeExpand<Self> {
32        self.into()
33    }
34}
35impl IntoExpand for e4m3 {
36    type Expand = NativeExpand<e4m3>;
37    fn into_expand(self, _scope: &Scope) -> Self::Expand {
38        self.into()
39    }
40}
41
42impl NativeAssign for e4m3 {}
43
44impl CubeType for e5m2 {
45    type ExpandType = NativeExpand<e5m2>;
46}
47
48impl CubeDebug for e5m2 {}
49impl Scalar for e5m2 {}
50impl CubePrimitive for e5m2 {
51    type Scalar = Self;
52    type Size = Const<1>;
53    type WithScalar<S: Scalar> = S;
54
55    /// Return the element type to use on GPU
56    fn as_type_native() -> Option<Type> {
57        Some(ElemType::Float(FloatKind::E5M2).into())
58    }
59
60    fn from_const_value(value: ConstantValue) -> Self {
61        let ConstantValue::Float(value) = value else {
62            unreachable!()
63        };
64        e5m2::from_f64(value)
65    }
66}
67
68impl IntoRuntime for e5m2 {
69    fn __expand_runtime_method(self, _scope: &Scope) -> NativeExpand<Self> {
70        self.into()
71    }
72}
73impl IntoExpand for e5m2 {
74    type Expand = NativeExpand<e5m2>;
75    fn into_expand(self, _scope: &Scope) -> Self::Expand {
76        self.into()
77    }
78}
79
80impl NativeAssign for e5m2 {}
81
82impl CubeType for ue8m0 {
83    type ExpandType = NativeExpand<ue8m0>;
84}
85
86impl CubeDebug for ue8m0 {}
87impl Scalar for ue8m0 {}
88impl CubePrimitive for ue8m0 {
89    type Scalar = Self;
90    type Size = Const<1>;
91    type WithScalar<S: Scalar> = S;
92
93    /// Return the element type to use on GPU
94    fn as_type_native() -> Option<Type> {
95        Some(ElemType::Float(FloatKind::UE8M0).into())
96    }
97
98    fn from_const_value(value: ConstantValue) -> Self {
99        let ConstantValue::Float(value) = value else {
100            unreachable!()
101        };
102        ue8m0::from_f64(value)
103    }
104}
105
106impl IntoRuntime for ue8m0 {
107    fn __expand_runtime_method(self, _scope: &Scope) -> NativeExpand<Self> {
108        self.into()
109    }
110}
111impl IntoExpand for ue8m0 {
112    type Expand = NativeExpand<ue8m0>;
113    fn into_expand(self, _scope: &Scope) -> Self::Expand {
114        self.into()
115    }
116}
117
118impl NativeAssign for ue8m0 {}