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 Scalar for e4m3 {}
11impl CubePrimitive for e4m3 {
12    type Scalar = Self;
13    type Size = Const<1>;
14    type WithScalar<S: Scalar> = S;
15
16    /// Return the element type to use on GPU
17    fn as_type_native() -> Option<Type> {
18        Some(ElemType::Float(FloatKind::E4M3).into())
19    }
20
21    fn from_const_value(value: ConstantValue) -> Self {
22        let ConstantValue::Float(value) = value else {
23            unreachable!()
24        };
25        e4m3::from_f64(value)
26    }
27}
28
29impl IntoRuntime for e4m3 {
30    fn __expand_runtime_method(self, _scope: &mut Scope) -> NativeExpand<Self> {
31        self.into()
32    }
33}
34
35impl NativeAssign for e4m3 {}
36
37impl CubeType for e5m2 {
38    type ExpandType = NativeExpand<e5m2>;
39}
40
41impl Scalar for e5m2 {}
42impl CubePrimitive for e5m2 {
43    type Scalar = Self;
44    type Size = Const<1>;
45    type WithScalar<S: Scalar> = S;
46
47    /// Return the element type to use on GPU
48    fn as_type_native() -> Option<Type> {
49        Some(ElemType::Float(FloatKind::E5M2).into())
50    }
51
52    fn from_const_value(value: ConstantValue) -> Self {
53        let ConstantValue::Float(value) = value else {
54            unreachable!()
55        };
56        e5m2::from_f64(value)
57    }
58}
59
60impl IntoRuntime for e5m2 {
61    fn __expand_runtime_method(self, _scope: &mut Scope) -> NativeExpand<Self> {
62        self.into()
63    }
64}
65
66impl NativeAssign for e5m2 {}
67
68impl CubeType for ue8m0 {
69    type ExpandType = NativeExpand<ue8m0>;
70}
71
72impl Scalar for ue8m0 {}
73impl CubePrimitive for ue8m0 {
74    type Scalar = Self;
75    type Size = Const<1>;
76    type WithScalar<S: Scalar> = S;
77
78    /// Return the element type to use on GPU
79    fn as_type_native() -> Option<Type> {
80        Some(ElemType::Float(FloatKind::UE8M0).into())
81    }
82
83    fn from_const_value(value: ConstantValue) -> Self {
84        let ConstantValue::Float(value) = value else {
85            unreachable!()
86        };
87        ue8m0::from_f64(value)
88    }
89}
90
91impl IntoRuntime for ue8m0 {
92    fn __expand_runtime_method(self, _scope: &mut Scope) -> NativeExpand<Self> {
93        self.into()
94    }
95}
96
97impl NativeAssign for ue8m0 {}