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

cubecl_core/frontend/element/
float.rs
