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