cubecl_common/float/
comptime_float.rs

1use core::fmt::{Debug, Display, Formatter};
2use core::hash::{Hash, Hasher};
3
4/// A float type whose bit representation can stand in for [`PartialEq`]/[`Eq`]/[`Hash`].
5///
6/// Implemented for the IEEE-754 types that need to be used as comptime kernel
7/// parameters: two values are equal iff their bit patterns are equal, so `-0.0`
8/// is normalized to `0.0` since they are numerically equal.
9pub trait FloatBits: Copy + PartialEq + Debug + Display {
10    /// The unsigned integer type with the same width as `Self`.
11    type Bits: Copy + Eq + Hash;
12
13    /// The bit pattern of this value.
14    fn to_bits(self) -> Self::Bits;
15
16    /// Whether this value is neither infinite nor NaN.
17    fn is_finite(self) -> bool;
18
19    /// Positive zero, used to normalize away the sign of zero.
20    fn zero() -> Self;
21}
22
23macro_rules! impl_float_bits {
24    ($ty:ty, $bits:ty) => {
25        impl FloatBits for $ty {
26            type Bits = $bits;
27
28            fn to_bits(self) -> $bits {
29                <$ty>::to_bits(self)
30            }
31
32            fn is_finite(self) -> bool {
33                <$ty>::is_finite(self)
34            }
35
36            fn zero() -> Self {
37                <$ty>::from_bits(0)
38            }
39        }
40    };
41}
42
43impl_float_bits!(f32, u32);
44impl_float_bits!(f64, u64);
45impl_float_bits!(half::f16, u16);
46impl_float_bits!(half::bf16, u16);
47
48/// A finite float usable as a comptime kernel parameter.
49///
50/// Plain floats can't back a [`Hash`]/[`Eq`] key because NaN breaks reflexivity,
51/// which is required for anything used as a kernel cache key (e.g. through
52/// `KernelId::info`). `ComptimeFloat` closes that gap by comparing and hashing
53/// bit patterns instead, after rejecting non-finite input and normalizing `-0.0`
54/// to `0.0` so numerically equal values always compare equal.
55///
56/// This is only appropriate for values that are genuinely arbitrary floats. If
57/// a value is always an exact ratio of two integers (e.g. derived from tensor
58/// shapes), prefer [`Ratio`](crate::Ratio): it is exact and avoids bit-pattern
59/// comparisons entirely.
60#[derive(Clone, Copy, Debug)]
61pub struct ComptimeFloat<F: FloatBits>(F);
62
63/// A float value that cannot be used as a [`ComptimeFloat`] because it is
64/// infinite or NaN.
65#[derive(Debug, Clone, Copy, PartialEq)]
66pub struct InvalidComptimeFloat<F>(pub F);
67
68impl<F: FloatBits> Display for InvalidComptimeFloat<F> {
69    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
70        write!(f, "invalid comptime float: {} is not finite", self.0)
71    }
72}
73
74impl<F: FloatBits> core::error::Error for InvalidComptimeFloat<F> {}
75
76impl<F: FloatBits> ComptimeFloat<F> {
77    /// Wrap `val`, rejecting non-finite input.
78    pub fn new(val: F) -> Result<Self, InvalidComptimeFloat<F>> {
79        if !val.is_finite() {
80            return Err(InvalidComptimeFloat(val));
81        }
82        // IEEE-754 equality treats -0.0 == 0.0, so this also normalizes -0.0.
83        let val = if val == F::zero() { F::zero() } else { val };
84        Ok(Self(val))
85    }
86
87    /// The wrapped value.
88    pub fn get(self) -> F {
89        self.0
90    }
91}
92
93impl<F: FloatBits> PartialEq for ComptimeFloat<F> {
94    fn eq(&self, other: &Self) -> bool {
95        self.0.to_bits() == other.0.to_bits()
96    }
97}
98
99impl<F: FloatBits> Eq for ComptimeFloat<F> {}
100
101impl<F: FloatBits> Hash for ComptimeFloat<F> {
102    fn hash<H: Hasher>(&self, state: &mut H) {
103        self.0.to_bits().hash(state);
104    }
105}
106
107impl<F: FloatBits> Display for ComptimeFloat<F> {
108    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
109        Display::fmt(&self.0, f)
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116    use alloc::format;
117
118    #[test]
119    fn rejects_non_finite() {
120        assert!(ComptimeFloat::new(f32::NAN).is_err());
121        assert!(ComptimeFloat::new(f32::INFINITY).is_err());
122        assert!(ComptimeFloat::new(f32::NEG_INFINITY).is_err());
123    }
124
125    #[test]
126    fn accepts_finite() {
127        assert!(ComptimeFloat::new(1.5f32).is_ok());
128    }
129
130    /// A trivial FNV-1a hasher so hash-equality can be tested without `std`.
131    #[derive(Default)]
132    struct TestHasher(u64);
133
134    impl Hasher for TestHasher {
135        fn finish(&self) -> u64 {
136            self.0
137        }
138
139        fn write(&mut self, bytes: &[u8]) {
140            self.0 = bytes.iter().fold(self.0, |hash, byte| {
141                (hash ^ *byte as u64).wrapping_mul(0x100000001b3)
142            });
143        }
144    }
145
146    fn hash_of<T: Hash>(val: &T) -> u64 {
147        let mut hasher = TestHasher::default();
148        val.hash(&mut hasher);
149        hasher.finish()
150    }
151
152    #[test]
153    fn negative_zero_equals_positive_zero() {
154        let neg = ComptimeFloat::new(-0.0f32).unwrap();
155        let pos = ComptimeFloat::new(0.0f32).unwrap();
156        assert_eq!(neg, pos);
157        assert_eq!(hash_of(&neg), hash_of(&pos));
158    }
159
160    #[test]
161    fn distinct_values_are_not_equal() {
162        let a = ComptimeFloat::new(1.0f32).unwrap();
163        let b = ComptimeFloat::new(2.0f32).unwrap();
164        assert_ne!(a, b);
165    }
166
167    #[test]
168    fn display_matches_inner_value() {
169        let val = ComptimeFloat::new(3.25f32).unwrap();
170        assert_eq!(format!("{}", val), "3.25");
171    }
172}
cubecl_common/float/comptime_float.rs

cubecl_common/float/
comptime_float.rs
