cubecl_core/frontend/element/
float.rs

1use cubecl_ir::{
2    ConstantValue, ElemType, Scope,
3    types::scalar::{BFloat16Type, Float16Type, Float32Type, Float64Type},
4};
5use half::{bf16, f16};
6use pliron::r#type::TypeHandle;
7
8use crate::{self as cubecl, ir::FloatKind, prelude::*};
9
10use super::Numeric;
11
12mod fp4;
13mod fp6;
14mod fp8;
15mod relaxed;
16mod tensor_float;
17
18/// Floating point numbers. Used as input in float kernels
19pub trait Float:
20    Numeric
21    + FloatOps
22    + ScalarNeg
23    + ScalarExp
24    + ScalarExp2
25    + ScalarLog
26    + ScalarLog1p
27    + ScalarExpm1
28    + ScalarCos
29    + ScalarSin
30    + ScalarTan
31    + ScalarTanh
32    + ScalarSinh
33    + ScalarCosh
34    + ScalarArcCos
35    + ScalarArcSin
36    + ScalarArcTan
37    + ScalarArcSinh
38    + ScalarArcCosh
39    + ScalarArcTanh
40    + ScalarDegrees
41    + ScalarRadians
42    + ScalarArcTan2
43    + ScalarPowf
44    + Powi<i32>
45    + ScalarHypot
46    + ScalarRhypot
47    + ScalarSqrt
48    + ScalarInverseSqrt
49    + ScalarCbrt
50    + ScalarInverseCbrt
51    + ScalarRound
52    + ScalarFloor
53    + ScalarCeil
54    + ScalarTrunc
55    + ScalarErf
56    + ScalarRecip
57    + ScalarMagnitude
58    + Normalize
59    + ScalarDot
60    + IsNan
61    + IsInf
62    + Into<Self::ExpandType>
63    + core::ops::Neg<Output = Self>
64    + core::cmp::PartialOrd
65    + core::cmp::PartialEq
66{
67    const DIGITS: u32;
68    const EPSILON: Self;
69    const INFINITY: Self;
70    const MANTISSA_DIGITS: u32;
71    const MAX_10_EXP: i32;
72    const MAX_EXP: i32;
73    const MIN_10_EXP: i32;
74    const MIN_EXP: i32;
75    const MIN_POSITIVE: Self;
76    const NAN: Self;
77    const NEG_INFINITY: Self;
78    const RADIX: u32;
79
80    fn new(val: f64) -> Self;
81    fn __expand_new(scope: &Scope, val: f64) -> <Self as CubeType>::ExpandType {
82        __expand_new(scope, val)
83    }
84}
85
86#[cube]
87pub trait FloatOps: CubePartialOrd + Sized {
88    fn min(self, other: Self) -> Self {
89        cubecl::prelude::min(self, other)
90    }
91
92    fn max(self, other: Self) -> Self {
93        cubecl::prelude::max(self, other)
94    }
95
96    fn clamp(self, min: Self, max: Self) -> Self {
97        clamp(self, min, max)
98    }
99}
100
101impl<T: Float> FloatOps for T {}
102impl<T: FloatOps + CubePrimitive> FloatOpsExpand for NativeExpand<T> {
103    fn __expand_min_method(self, scope: &Scope, other: Self) -> Self {
104        min::expand(scope, self, other)
105    }
106
107    fn __expand_max_method(self, scope: &Scope, other: Self) -> Self {
108        max::expand(scope, self, other)
109    }
110
111    fn __expand_clamp_method(self, scope: &Scope, min: Self, max: Self) -> Self {
112        clamp::expand(scope, self, min, max)
113    }
114}
115
116macro_rules! impl_float {
117    (half $primitive:ident, $ty: ty, $kind:ident) => {
118        impl_float!($primitive, $ty, $kind, |val| $primitive::from_f64(val));
119    };
120    ($primitive:ident, $ty: ty, $kind:ident) => {
121        impl_float!($primitive, $ty, $kind, |val| val as $primitive);
122    };
123    ($primitive:ident, $ty: ty, $kind:ident, $new:expr) => {
124        impl CubeType for $primitive {
125            type ExpandType = NativeExpand<$primitive>;
126        }
127
128        impl CubeDebug for $primitive {}
129        impl Scalar for $primitive {
130            fn elem_type_native() -> ElemType {
131                FloatKind::$kind.into()
132            }
133        }
134        impl CubePrimitive for $primitive {
135            type Scalar = Self;
136            type Size = Const<1>;
137            type WithScalar<S: Scalar> = S;
138
139            fn from_const_value(value: ConstantValue) -> Self {
140                let ConstantValue::Float(value) = value else {
141                    unreachable!()
142                };
143                $new(value)
144            }
145        }
146
147        impl ConstantTyped for $primitive {
148            fn __expand_as_type(scope: &Scope) -> TypeHandle {
149                <$ty>::get(scope.ctx()).into()
150            }
151        }
152
153        impl IntoRuntime for $primitive {
154            fn __expand_runtime_method(self, _scope: &Scope) -> NativeExpand<Self> {
155                self.into()
156            }
157        }
158
159        impl IntoExpand for $primitive {
160            type Expand = NativeExpand<$primitive>;
161
162            fn into_expand(self, _: &Scope) -> Self::Expand {
163                self.into()
164            }
165        }
166
167        impl Numeric for $primitive {
168            fn min_value() -> Self {
169                <Self as num_traits::Float>::min_value()
170            }
171            fn max_value() -> Self {
172                <Self as num_traits::Float>::max_value()
173            }
174        }
175
176        impl NativeAssign for $primitive {}
177
178        impl IntoMut for $primitive {
179            fn into_mut(self, _scope: &Scope) -> Self {
180                self
181            }
182        }
183
184        impl Float for $primitive {
185            const DIGITS: u32 = $primitive::DIGITS;
186            const EPSILON: Self = $primitive::EPSILON;
187            const INFINITY: Self = $primitive::INFINITY;
188            const MANTISSA_DIGITS: u32 = $primitive::MANTISSA_DIGITS;
189            const MAX_10_EXP: i32 = $primitive::MAX_10_EXP;
190            const MAX_EXP: i32 = $primitive::MAX_EXP;
191            const MIN_10_EXP: i32 = $primitive::MIN_10_EXP;
192            const MIN_EXP: i32 = $primitive::MIN_EXP;
193            const MIN_POSITIVE: Self = $primitive::MIN_POSITIVE;
194            const NAN: Self = $primitive::NAN;
195            const NEG_INFINITY: Self = $primitive::NEG_INFINITY;
196            const RADIX: u32 = $primitive::RADIX;
197
198            fn new(val: f64) -> Self {
199                $new(val)
200            }
201        }
202
203        impl_scalar_launch!($primitive);
204    };
205}
206
207impl_float!(half f16, Float16Type, F16);
208impl_float!(half bf16,  BFloat16Type, BF16);
209impl_float!(f32, Float32Type, F32);
210impl_float!(f64, Float64Type, F64);
cubecl_core/frontend/element/float.rs

cubecl_core/frontend/element/
float.rs
