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 __expand_as_type(scope: &Scope) -> TypeHandle {
26        Float8E4M3Type::get(scope.ctx()).into()
27    }
28
29    fn from_const_value(value: ConstantValue) -> Self {
30        let ConstantValue::Float(value) = value else {
31            unreachable!()
32        };
33        e4m3::from_f64(value)
34    }
35}
36
37impl IntoRuntime for e4m3 {
38    fn __expand_runtime_method(self, _scope: &Scope) -> NativeExpand<Self> {
39        self.into()
40    }
41}
42impl IntoExpand for e4m3 {
43    type Expand = NativeExpand<e4m3>;
44    fn into_expand(self, _scope: &Scope) -> Self::Expand {
45        self.into()
46    }
47}
48
49impl NativeAssign for e4m3 {}
50
51impl CubeType for e5m2 {
52    type ExpandType = NativeExpand<e5m2>;
53}
54
55impl CubeDebug for e5m2 {}
56impl Scalar for e5m2 {
57    fn elem_type_native() -> ElemType {
58        FloatKind::E5M2.into()
59    }
60}
61impl CubePrimitive for e5m2 {
62    type Scalar = Self;
63    type Size = Const<1>;
64    type WithScalar<S: Scalar> = S;
65
66    fn __expand_as_type(scope: &Scope) -> TypeHandle {
67        Float8E5M2Type::get(scope.ctx()).into()
68    }
69
70    fn from_const_value(value: ConstantValue) -> Self {
71        let ConstantValue::Float(value) = value else {
72            unreachable!()
73        };
74        e5m2::from_f64(value)
75    }
76}
77
78impl IntoRuntime for e5m2 {
79    fn __expand_runtime_method(self, _scope: &Scope) -> NativeExpand<Self> {
80        self.into()
81    }
82}
83impl IntoExpand for e5m2 {
84    type Expand = NativeExpand<e5m2>;
85    fn into_expand(self, _scope: &Scope) -> Self::Expand {
86        self.into()
87    }
88}
89
90impl NativeAssign for e5m2 {}
91
92impl CubeType for ue8m0 {
93    type ExpandType = NativeExpand<ue8m0>;
94}
95
96impl CubeDebug for ue8m0 {}
97impl Scalar for ue8m0 {
98    fn elem_type_native() -> ElemType {
99        FloatKind::UE8M0.into()
100    }
101}
102impl CubePrimitive for ue8m0 {
103    type Scalar = Self;
104    type Size = Const<1>;
105    type WithScalar<S: Scalar> = S;
106
107    fn __expand_as_type(scope: &Scope) -> TypeHandle {
108        Float8E8M0Type::get(scope.ctx()).into()
109    }
110
111    fn from_const_value(value: ConstantValue) -> Self {
112        let ConstantValue::Float(value) = value else {
113            unreachable!()
114        };
115        ue8m0::from_f64(value)
116    }
117}
118
119impl IntoRuntime for ue8m0 {
120    fn __expand_runtime_method(self, _scope: &Scope) -> NativeExpand<Self> {
121        self.into()
122    }
123}
124impl IntoExpand for ue8m0 {
125    type Expand = NativeExpand<ue8m0>;
126    fn into_expand(self, _scope: &Scope) -> Self::Expand {
127        self.into()
128    }
129}
130
131impl NativeAssign for ue8m0 {}
cubecl_core/frontend/element/float/fp8.rs

cubecl_core/frontend/element/float/
fp8.rs
