cubecl_core/frontend/element/float/
fp8.rs

1use cubecl_common::{e4m3, e5m2, ue8m0};
2use cubecl_ir::{Elem, ExpandElement, FloatKind, Scope};
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_elem_native() -> Option<Elem> {
16        Some(Elem::Float(FloatKind::E4M3))
17    }
18}
19
20impl IntoRuntime for e4m3 {
21    fn __expand_runtime_method(self, scope: &mut Scope) -> ExpandElementTyped<Self> {
22        let elem: ExpandElementTyped<Self> = self.into();
23        into_runtime_expand_element(scope, elem).into()
24    }
25}
26
27impl ExpandElementIntoMut for e4m3 {
28    fn elem_into_mut(scope: &mut Scope, elem: ExpandElement) -> ExpandElement {
29        into_mut_expand_element(scope, elem)
30    }
31}
32
33impl CubeType for e5m2 {
34    type ExpandType = ExpandElementTyped<e5m2>;
35}
36
37impl CubePrimitive for e5m2 {
38    /// Return the element type to use on GPU
39    fn as_elem_native() -> Option<Elem> {
40        Some(Elem::Float(FloatKind::E5M2))
41    }
42}
43
44impl IntoRuntime for e5m2 {
45    fn __expand_runtime_method(self, scope: &mut Scope) -> ExpandElementTyped<Self> {
46        let elem: ExpandElementTyped<Self> = self.into();
47        into_runtime_expand_element(scope, elem).into()
48    }
49}
50
51impl ExpandElementIntoMut for e5m2 {
52    fn elem_into_mut(scope: &mut Scope, elem: ExpandElement) -> ExpandElement {
53        into_mut_expand_element(scope, elem)
54    }
55}
56
57impl CubeType for ue8m0 {
58    type ExpandType = ExpandElementTyped<ue8m0>;
59}
60
61impl CubePrimitive for ue8m0 {
62    /// Return the element type to use on GPU
63    fn as_elem_native() -> Option<Elem> {
64        Some(Elem::Float(FloatKind::UE8M0))
65    }
66}
67
68impl IntoRuntime for ue8m0 {
69    fn __expand_runtime_method(self, scope: &mut Scope) -> ExpandElementTyped<Self> {
70        let elem: ExpandElementTyped<Self> = self.into();
71        into_runtime_expand_element(scope, elem).into()
72    }
73}
74
75impl ExpandElementIntoMut for ue8m0 {
76    fn elem_into_mut(scope: &mut Scope, elem: ExpandElement) -> ExpandElement {
77        into_mut_expand_element(scope, elem)
78    }
79}
cubecl_core/frontend/element/float/fp8.rs

cubecl_core/frontend/element/float/
fp8.rs
