cubecl_core/frontend/element/float/
fp4.rs

1use cubecl_common::{e2m1, e2m1x2};
2use cubecl_ir::{ConstantValue, ElemType, FloatKind, Scope, StorageType, Type};
3
4use crate::prelude::*;
5
6impl CubeType for e2m1 {
7    type ExpandType = NativeExpand<e2m1>;
8}
9
10impl CubeDebug for e2m1 {}
11impl Scalar for e2m1 {}
12impl CubePrimitive for e2m1 {
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(StorageType::Scalar(ElemType::Float(FloatKind::E2M1)).into())
20    }
21
22    fn from_const_value(value: ConstantValue) -> Self {
23        let ConstantValue::Float(value) = value else {
24            unreachable!()
25        };
26        e2m1::from_f64(value)
27    }
28}
29
30impl IntoRuntime for e2m1 {
31    fn __expand_runtime_method(self, _scope: &Scope) -> NativeExpand<Self> {
32        self.into()
33    }
34}
35impl IntoExpand for e2m1 {
36    type Expand = NativeExpand<e2m1>;
37    fn into_expand(self, _scope: &Scope) -> Self::Expand {
38        self.into()
39    }
40}
41
42impl NativeAssign for e2m1 {}
43
44impl CubeType for e2m1x2 {
45    type ExpandType = NativeExpand<e2m1x2>;
46}
47
48impl CubeDebug for e2m1x2 {}
49// Considered a scalar because it's really just a `u8` in a trenchcoat, and should be possible to
50// store in a `Vector`.
51impl Scalar for e2m1x2 {}
52impl CubePrimitive for e2m1x2 {
53    type Scalar = Self;
54    type Size = Const<1>;
55    type WithScalar<S: Scalar> = S;
56
57    /// Return the element type to use on GPU
58    fn as_type_native() -> Option<Type> {
59        Some(StorageType::Packed(ElemType::Float(FloatKind::E2M1), 2).into())
60    }
61
62    fn from_const_value(value: ConstantValue) -> Self {
63        let ConstantValue::Float(value) = value else {
64            unreachable!()
65        };
66        let val = e2m1::from_f64(value).to_bits();
67        // Fill both values, not sure this is ever useful but it works
68        e2m1x2::from_bits(val | (val << 4))
69    }
70}
71
72impl IntoRuntime for e2m1x2 {
73    fn __expand_runtime_method(self, _scope: &Scope) -> NativeExpand<Self> {
74        self.into()
75    }
76}
77impl IntoExpand for e2m1x2 {
78    type Expand = NativeExpand<e2m1x2>;
79    fn into_expand(self, _scope: &Scope) -> Self::Expand {
80        self.into()
81    }
82}
83
84impl NativeAssign for e2m1x2 {}
cubecl_core/frontend/element/float/fp4.rs

cubecl_core/frontend/element/float/
fp4.rs
