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

cubecl_core/frontend/element/
float.rs
