Skip to main content

cubecl_core/frontend/element/float/
fp8.rs

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