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

cubecl_core/frontend/element/float/
fp4.rs
