ruda_kernel/dsl/frontend/element/
float.rs

1use ruda_core::ir::{ConstantValue, Scope, StorageType, Type};
2use half::{bf16, f16};
3
4use crate::dsl::{
5    ir::{ElemType, FloatKind},
6    prelude::*,
7};
8
9use super::Numeric;
10
11mod fp4;
12mod fp6;
13mod fp8;
14mod relaxed;
15mod tensor_float;
16
17/// Floating point numbers. Used as input in float kernels
18pub trait Float:
19    Numeric
20    + FloatOps
21    + Exp
22    + Log
23    + Log1p
24    + Cos
25    + Sin
26    + Tan
27    + Tanh
28    + Sinh
29    + Cosh
30    + ArcCos
31    + ArcSin
32    + ArcTan
33    + ArcSinh
34    + ArcCosh
35    + ArcTanh
36    + Degrees
37    + Radians
38    + ArcTan2
39    + Powf
40    + Powi<i32>
41    + Hypot
42    + Rhypot
43    + Sqrt
44    + InverseSqrt
45    + Round
46    + Floor
47    + Ceil
48    + Trunc
49    + Erf
50    + Recip
51    + Magnitude
52    + Normalize
53    + Dot
54    + IsNan
55    + IsInf
56    + Into<Self::ExpandType>
57    + core::ops::Neg<Output = Self>
58    + core::cmp::PartialOrd
59    + core::cmp::PartialEq
60{
61    const DIGITS: u32;
62    const EPSILON: Self;
63    const INFINITY: Self;
64    const MANTISSA_DIGITS: u32;
65    const MAX_10_EXP: i32;
66    const MAX_EXP: i32;
67    const MIN_10_EXP: i32;
68    const MIN_EXP: i32;
69    const MIN_POSITIVE: Self;
70    const NAN: Self;
71    const NEG_INFINITY: Self;
72    const RADIX: u32;
73
74    fn new(val: f32) -> Self;
75    fn __expand_new(scope: &mut Scope, val: f32) -> <Self as RudaType>::ExpandType {
76        __expand_new(scope, val)
77    }
78}
79
80#[ruda]
81pub trait FloatOps: RudaPrimitive + PartialOrd + Sized {
82    fn min(self, other: Self) -> Self {
83        crate::dsl::prelude::min(self, other)
84    }
85
86    fn max(self, other: Self) -> Self {
87        crate::dsl::prelude::max(self, other)
88    }
89
90    fn clamp(self, min: Self, max: Self) -> Self {
91        clamp(self, min, max)
92    }
93}
94
95impl<T: Float> FloatOps for T {}
96impl<T: FloatOps + RudaPrimitive> FloatOpsExpand for NativeExpand<T> {
97    fn __expand_min_method(self, scope: &mut Scope, other: Self) -> Self {
98        min::expand(scope, self, other)
99    }
100
101    fn __expand_max_method(self, scope: &mut Scope, other: Self) -> Self {
102        max::expand(scope, self, other)
103    }
104
105    fn __expand_clamp_method(self, scope: &mut Scope, min: Self, max: Self) -> Self {
106        clamp::expand(scope, self, min, max)
107    }
108}
109
110macro_rules! impl_float {
111    (half $primitive:ident, $kind:ident) => {
112        impl_float!($primitive, $kind, |val| $primitive::from_f64(val));
113    };
114    ($primitive:ident, $kind:ident) => {
115        impl_float!($primitive, $kind, |val| val as $primitive);
116    };
117    ($primitive:ident, $kind:ident, $new:expr) => {
118        impl RudaType for $primitive {
119            type ExpandType = NativeExpand<$primitive>;
120        }
121
122        impl Scalar for $primitive {}
123        impl RudaPrimitive for $primitive {
124            type Scalar = Self;
125            type Size = Const<1>;
126            type WithScalar<S: Scalar> = S;
127
128            /// Return the element type to use on GPU
129            fn as_type_native() -> Option<Type> {
130                Some(StorageType::Scalar(ElemType::Float(FloatKind::$kind)).into())
131            }
132
133            fn from_const_value(value: ConstantValue) -> Self {
134                let ConstantValue::Float(value) = value else {
135                    unreachable!()
136                };
137                $new(value)
138            }
139        }
140
141        impl IntoRuntime for $primitive {
142            fn __expand_runtime_method(self, _scope: &mut Scope) -> NativeExpand<Self> {
143                self.into()
144            }
145        }
146
147        impl Numeric for $primitive {
148            fn min_value() -> Self {
149                <Self as num_traits::Float>::min_value()
150            }
151            fn max_value() -> Self {
152                <Self as num_traits::Float>::max_value()
153            }
154        }
155
156        impl NativeAssign for $primitive {}
157
158        impl IntoMut for $primitive {
159            fn into_mut(self, _scope: &mut Scope) -> Self {
160                self
161            }
162        }
163
164        impl Float for $primitive {
165            const DIGITS: u32 = $primitive::DIGITS;
166            const EPSILON: Self = $primitive::EPSILON;
167            const INFINITY: Self = $primitive::INFINITY;
168            const MANTISSA_DIGITS: u32 = $primitive::MANTISSA_DIGITS;
169            const MAX_10_EXP: i32 = $primitive::MAX_10_EXP;
170            const MAX_EXP: i32 = $primitive::MAX_EXP;
171            const MIN_10_EXP: i32 = $primitive::MIN_10_EXP;
172            const MIN_EXP: i32 = $primitive::MIN_EXP;
173            const MIN_POSITIVE: Self = $primitive::MIN_POSITIVE;
174            const NAN: Self = $primitive::NAN;
175            const NEG_INFINITY: Self = $primitive::NEG_INFINITY;
176            const RADIX: u32 = $primitive::RADIX;
177
178            fn new(val: f32) -> Self {
179                $new(val as f64)
180            }
181        }
182    };
183}
184
185impl_float!(half f16, F16);
186impl_float!(half bf16, BF16);
187impl_float!(f32, F32);
188impl_float!(f64, F64);
ruda_kernel/dsl/frontend/element/float.rs

ruda_kernel/dsl/frontend/element/
float.rs
