cubecl_core/frontend/element/float/
fp4.rs

1use cubecl_common::{e2m1, e2m1x2};
2use cubecl_ir::{
3    ConstantValue, ElemType, FloatKind, Scope,
4    types::scalar::{Float4E2M1Type, Float4E2M1x2Type},
5};
6use pliron::r#type::TypeHandle;
7
8use crate::prelude::*;
9
10impl CubeType for e2m1 {
11    type ExpandType = NativeExpand<e2m1>;
12}
13
14impl CubeDebug for e2m1 {}
15impl Scalar for e2m1 {
16    fn elem_type_native() -> ElemType {
17        FloatKind::E2M1.into()
18    }
19}
20impl CubePrimitive for e2m1 {
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        Float4E2M1Type::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        e2m1::from_f64(value)
34    }
35}
36
37impl IntoRuntime for e2m1 {
38    fn __expand_runtime_method(self, _scope: &Scope) -> NativeExpand<Self> {
39        self.into()
40    }
41}
42impl IntoExpand for e2m1 {
43    type Expand = NativeExpand<e2m1>;
44    fn into_expand(self, _scope: &Scope) -> Self::Expand {
45        self.into()
46    }
47}
48
49impl NativeAssign for e2m1 {}
50
51impl CubeType for e2m1x2 {
52    type ExpandType = NativeExpand<e2m1x2>;
53}
54
55impl CubeDebug for e2m1x2 {}
56// Considered a scalar because it's really just a `u8` in a trenchcoat, and should be possible to
57// store in a `Vector`.
58impl Scalar for e2m1x2 {
59    fn elem_type_native() -> ElemType {
60        FloatKind::E2M1x2.into()
61    }
62}
63impl CubePrimitive for e2m1x2 {
64    type Scalar = Self;
65    type Size = Const<1>;
66    type WithScalar<S: Scalar> = S;
67
68    fn __expand_as_type(scope: &Scope) -> TypeHandle {
69        Float4E2M1x2Type::get(scope.ctx()).into()
70    }
71
72    fn from_const_value(value: ConstantValue) -> Self {
73        let ConstantValue::Float(value) = value else {
74            unreachable!()
75        };
76        let val = e2m1::from_f64(value).to_bits();
77        // Fill both values, not sure this is ever useful but it works
78        e2m1x2::from_bits(val | (val << 4))
79    }
80}
81
82impl IntoRuntime for e2m1x2 {
83    fn __expand_runtime_method(self, _scope: &Scope) -> NativeExpand<Self> {
84        self.into()
85    }
86}
87impl IntoExpand for e2m1x2 {
88    type Expand = NativeExpand<e2m1x2>;
89    fn into_expand(self, _scope: &Scope) -> Self::Expand {
90        self.into()
91    }
92}
93
94impl NativeAssign for e2m1x2 {}
cubecl_core/frontend/element/float/fp4.rs

cubecl_core/frontend/element/float/
fp4.rs
