• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 //! Contains utility functions and traits to convert between vectors of [`u16`] bits and [`f16`] or
2 //! [`bf16`] vectors.
3 //!
4 //! The utility [`HalfBitsVecExt`] sealed extension trait is implemented for [`Vec<u16>`] vectors,
5 //! while the utility [`HalfFloatVecExt`] sealed extension trait is implemented for both
6 //! [`Vec<f16>`] and [`Vec<bf16>`] vectors. These traits provide efficient conversions and
7 //! reinterpret casting of larger buffers of floating point values, and are automatically included
8 //! in the [`prelude`][crate::prelude] module.
9 //!
10 //! This module is only available with the `std` or `alloc` feature.
11 
12 use super::{bf16, f16, slice::HalfFloatSliceExt};
13 #[cfg(feature = "alloc")]
14 use alloc::vec::Vec;
15 use core::mem;
16 
17 /// Extensions to [`Vec<f16>`] and [`Vec<bf16>`] to support reinterpret operations.
18 ///
19 /// This trait is sealed and cannot be implemented outside of this crate.
20 pub trait HalfFloatVecExt: private::SealedHalfFloatVec {
21     /// Reinterprets a vector of [`f16`]or [`bf16`] numbers as a vector of [`u16`] bits.
22     ///
23     /// This is a zero-copy operation. The reinterpreted vector has the same memory location as
24     /// `self`.
25     ///
26     /// # Examples
27     ///
28     /// ```rust
29     /// # use half::prelude::*;
30     /// let float_buffer = vec![f16::from_f32(1.), f16::from_f32(2.), f16::from_f32(3.)];
31     /// let int_buffer = float_buffer.reinterpret_into();
32     ///
33     /// assert_eq!(int_buffer, [f16::from_f32(1.).to_bits(), f16::from_f32(2.).to_bits(), f16::from_f32(3.).to_bits()]);
34     /// ```
35     #[must_use]
reinterpret_into(self) -> Vec<u16>36     fn reinterpret_into(self) -> Vec<u16>;
37 
38     /// Converts all of the elements of a `[f32]` slice into a new [`f16`] or [`bf16`] vector.
39     ///
40     /// The conversion operation is vectorized over the slice, meaning the conversion may be more
41     /// efficient than converting individual elements on some hardware that supports SIMD
42     /// conversions. See [crate documentation][crate] for more information on hardware conversion
43     /// support.
44     ///
45     /// # Examples
46     /// ```rust
47     /// # use half::prelude::*;
48     /// let float_values = [1., 2., 3., 4.];
49     /// let vec: Vec<f16> = Vec::from_f32_slice(&float_values);
50     ///
51     /// assert_eq!(vec, vec![f16::from_f32(1.), f16::from_f32(2.), f16::from_f32(3.), f16::from_f32(4.)]);
52     /// ```
53     #[must_use]
from_f32_slice(slice: &[f32]) -> Self54     fn from_f32_slice(slice: &[f32]) -> Self;
55 
56     /// Converts all of the elements of a `[f64]` slice into a new [`f16`] or [`bf16`] vector.
57     ///
58     /// The conversion operation is vectorized over the slice, meaning the conversion may be more
59     /// efficient than converting individual elements on some hardware that supports SIMD
60     /// conversions. See [crate documentation][crate] for more information on hardware conversion
61     /// support.
62     ///
63     /// # Examples
64     /// ```rust
65     /// # use half::prelude::*;
66     /// let float_values = [1., 2., 3., 4.];
67     /// let vec: Vec<f16> = Vec::from_f64_slice(&float_values);
68     ///
69     /// assert_eq!(vec, vec![f16::from_f64(1.), f16::from_f64(2.), f16::from_f64(3.), f16::from_f64(4.)]);
70     /// ```
71     #[must_use]
from_f64_slice(slice: &[f64]) -> Self72     fn from_f64_slice(slice: &[f64]) -> Self;
73 }
74 
75 /// Extensions to [`Vec<u16>`] to support reinterpret operations.
76 ///
77 /// This trait is sealed and cannot be implemented outside of this crate.
78 pub trait HalfBitsVecExt: private::SealedHalfBitsVec {
79     /// Reinterprets a vector of [`u16`] bits as a vector of [`f16`] or [`bf16`] numbers.
80     ///
81     /// `H` is the type to cast to, and must be either the [`f16`] or [`bf16`] type.
82     ///
83     /// This is a zero-copy operation. The reinterpreted vector has the same memory location as
84     /// `self`.
85     ///
86     /// # Examples
87     ///
88     /// ```rust
89     /// # use half::prelude::*;
90     /// let int_buffer = vec![f16::from_f32(1.).to_bits(), f16::from_f32(2.).to_bits(), f16::from_f32(3.).to_bits()];
91     /// let float_buffer = int_buffer.reinterpret_into::<f16>();
92     ///
93     /// assert_eq!(float_buffer, [f16::from_f32(1.), f16::from_f32(2.), f16::from_f32(3.)]);
94     /// ```
95     #[must_use]
reinterpret_into<H>(self) -> Vec<H> where H: crate::private::SealedHalf96     fn reinterpret_into<H>(self) -> Vec<H>
97     where
98         H: crate::private::SealedHalf;
99 }
100 
101 mod private {
102     use crate::{bf16, f16};
103     #[cfg(feature = "alloc")]
104     use alloc::vec::Vec;
105 
106     pub trait SealedHalfFloatVec {}
107     impl SealedHalfFloatVec for Vec<f16> {}
108     impl SealedHalfFloatVec for Vec<bf16> {}
109 
110     pub trait SealedHalfBitsVec {}
111     impl SealedHalfBitsVec for Vec<u16> {}
112 }
113 
114 impl HalfFloatVecExt for Vec<f16> {
115     #[inline]
reinterpret_into(mut self) -> Vec<u16>116     fn reinterpret_into(mut self) -> Vec<u16> {
117         // An f16 array has same length and capacity as u16 array
118         let length = self.len();
119         let capacity = self.capacity();
120 
121         // Actually reinterpret the contents of the Vec<f16> as u16,
122         // knowing that structs are represented as only their members in memory,
123         // which is the u16 part of `f16(u16)`
124         let pointer = self.as_mut_ptr() as *mut u16;
125 
126         // Prevent running a destructor on the old Vec<u16>, so the pointer won't be deleted
127         mem::forget(self);
128 
129         // Finally construct a new Vec<f16> from the raw pointer
130         // SAFETY: We are reconstructing full length and capacity of original vector,
131         // using its original pointer, and the size of elements are identical.
132         unsafe { Vec::from_raw_parts(pointer, length, capacity) }
133     }
134 
135     #[allow(clippy::uninit_vec)]
from_f32_slice(slice: &[f32]) -> Self136     fn from_f32_slice(slice: &[f32]) -> Self {
137         let mut vec = Vec::with_capacity(slice.len());
138         // SAFETY: convert will initialize every value in the vector without reading them,
139         // so this is safe to do instead of double initialize from resize, and we're setting it to
140         // same value as capacity.
141         unsafe { vec.set_len(slice.len()) };
142         vec.convert_from_f32_slice(slice);
143         vec
144     }
145 
146     #[allow(clippy::uninit_vec)]
from_f64_slice(slice: &[f64]) -> Self147     fn from_f64_slice(slice: &[f64]) -> Self {
148         let mut vec = Vec::with_capacity(slice.len());
149         // SAFETY: convert will initialize every value in the vector without reading them,
150         // so this is safe to do instead of double initialize from resize, and we're setting it to
151         // same value as capacity.
152         unsafe { vec.set_len(slice.len()) };
153         vec.convert_from_f64_slice(slice);
154         vec
155     }
156 }
157 
158 impl HalfFloatVecExt for Vec<bf16> {
159     #[inline]
reinterpret_into(mut self) -> Vec<u16>160     fn reinterpret_into(mut self) -> Vec<u16> {
161         // An f16 array has same length and capacity as u16 array
162         let length = self.len();
163         let capacity = self.capacity();
164 
165         // Actually reinterpret the contents of the Vec<f16> as u16,
166         // knowing that structs are represented as only their members in memory,
167         // which is the u16 part of `f16(u16)`
168         let pointer = self.as_mut_ptr() as *mut u16;
169 
170         // Prevent running a destructor on the old Vec<u16>, so the pointer won't be deleted
171         mem::forget(self);
172 
173         // Finally construct a new Vec<f16> from the raw pointer
174         // SAFETY: We are reconstructing full length and capacity of original vector,
175         // using its original pointer, and the size of elements are identical.
176         unsafe { Vec::from_raw_parts(pointer, length, capacity) }
177     }
178 
179     #[allow(clippy::uninit_vec)]
from_f32_slice(slice: &[f32]) -> Self180     fn from_f32_slice(slice: &[f32]) -> Self {
181         let mut vec = Vec::with_capacity(slice.len());
182         // SAFETY: convert will initialize every value in the vector without reading them,
183         // so this is safe to do instead of double initialize from resize, and we're setting it to
184         // same value as capacity.
185         unsafe { vec.set_len(slice.len()) };
186         vec.convert_from_f32_slice(slice);
187         vec
188     }
189 
190     #[allow(clippy::uninit_vec)]
from_f64_slice(slice: &[f64]) -> Self191     fn from_f64_slice(slice: &[f64]) -> Self {
192         let mut vec = Vec::with_capacity(slice.len());
193         // SAFETY: convert will initialize every value in the vector without reading them,
194         // so this is safe to do instead of double initialize from resize, and we're setting it to
195         // same value as capacity.
196         unsafe { vec.set_len(slice.len()) };
197         vec.convert_from_f64_slice(slice);
198         vec
199     }
200 }
201 
202 impl HalfBitsVecExt for Vec<u16> {
203     // This is safe because all traits are sealed
204     #[inline]
reinterpret_into<H>(mut self) -> Vec<H> where H: crate::private::SealedHalf,205     fn reinterpret_into<H>(mut self) -> Vec<H>
206     where
207         H: crate::private::SealedHalf,
208     {
209         // An f16 array has same length and capacity as u16 array
210         let length = self.len();
211         let capacity = self.capacity();
212 
213         // Actually reinterpret the contents of the Vec<u16> as f16,
214         // knowing that structs are represented as only their members in memory,
215         // which is the u16 part of `f16(u16)`
216         let pointer = self.as_mut_ptr() as *mut H;
217 
218         // Prevent running a destructor on the old Vec<u16>, so the pointer won't be deleted
219         mem::forget(self);
220 
221         // Finally construct a new Vec<f16> from the raw pointer
222         // SAFETY: We are reconstructing full length and capacity of original vector,
223         // using its original pointer, and the size of elements are identical.
224         unsafe { Vec::from_raw_parts(pointer, length, capacity) }
225     }
226 }
227 
228 #[cfg(test)]
229 mod test {
230     use super::{HalfBitsVecExt, HalfFloatVecExt};
231     use crate::{bf16, f16};
232     #[cfg(all(feature = "alloc", not(feature = "std")))]
233     use alloc::vec;
234 
235     #[test]
test_vec_conversions_f16()236     fn test_vec_conversions_f16() {
237         let numbers = vec![f16::E, f16::PI, f16::EPSILON, f16::FRAC_1_SQRT_2];
238         let bits = vec![
239             f16::E.to_bits(),
240             f16::PI.to_bits(),
241             f16::EPSILON.to_bits(),
242             f16::FRAC_1_SQRT_2.to_bits(),
243         ];
244         let bits_cloned = bits.clone();
245 
246         // Convert from bits to numbers
247         let from_bits = bits.reinterpret_into::<f16>();
248         assert_eq!(&from_bits[..], &numbers[..]);
249 
250         // Convert from numbers back to bits
251         let to_bits = from_bits.reinterpret_into();
252         assert_eq!(&to_bits[..], &bits_cloned[..]);
253     }
254 
255     #[test]
test_vec_conversions_bf16()256     fn test_vec_conversions_bf16() {
257         let numbers = vec![bf16::E, bf16::PI, bf16::EPSILON, bf16::FRAC_1_SQRT_2];
258         let bits = vec![
259             bf16::E.to_bits(),
260             bf16::PI.to_bits(),
261             bf16::EPSILON.to_bits(),
262             bf16::FRAC_1_SQRT_2.to_bits(),
263         ];
264         let bits_cloned = bits.clone();
265 
266         // Convert from bits to numbers
267         let from_bits = bits.reinterpret_into::<bf16>();
268         assert_eq!(&from_bits[..], &numbers[..]);
269 
270         // Convert from numbers back to bits
271         let to_bits = from_bits.reinterpret_into();
272         assert_eq!(&to_bits[..], &bits_cloned[..]);
273     }
274 }
275