260 lines
9.5 KiB
Rust
260 lines
9.5 KiB
Rust
//! Contains utility functions and traits to convert between vectors of [`u16`] bits and [`struct@f16`] or
|
|
//! [`bf16`] vectors.
|
|
//!
|
|
//! The utility [`HalfBitsVecExt`] sealed extension trait is implemented for [`Vec<u16>`] vectors,
|
|
//! while the utility [`HalfFloatVecExt`] sealed extension trait is implemented for both
|
|
//! [`Vec<f16>`] and [`Vec<bf16>`] vectors. These traits provide efficient conversions and
|
|
//! reinterpret casting of larger buffers of floating point values, and are automatically included
|
|
//! in the [`prelude`][crate::prelude] module.
|
|
//!
|
|
//! This module is only available with the `std` or `alloc` feature.
|
|
|
|
use super::{bf16, f16, slice::HalfFloatSliceExt};
|
|
#[cfg(feature = "alloc")]
|
|
#[allow(unused_imports)]
|
|
use alloc::{vec, vec::Vec};
|
|
use core::mem;
|
|
|
|
/// Extensions to [`Vec<f16>`] and [`Vec<bf16>`] to support reinterpret operations.
|
|
///
|
|
/// This trait is sealed and cannot be implemented outside of this crate.
|
|
pub trait HalfFloatVecExt: private::SealedHalfFloatVec {
|
|
/// Reinterprets a vector of [`struct@f16`]or [`bf16`] numbers as a vector of [`u16`] bits.
|
|
///
|
|
/// This is a zero-copy operation. The reinterpreted vector has the same memory location as
|
|
/// `self`.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```rust
|
|
/// # use half::prelude::*;
|
|
/// let float_buffer = vec![f16::from_f32(1.), f16::from_f32(2.), f16::from_f32(3.)];
|
|
/// let int_buffer = float_buffer.reinterpret_into();
|
|
///
|
|
/// assert_eq!(int_buffer, [f16::from_f32(1.).to_bits(), f16::from_f32(2.).to_bits(), f16::from_f32(3.).to_bits()]);
|
|
/// ```
|
|
#[must_use]
|
|
fn reinterpret_into(self) -> Vec<u16>;
|
|
|
|
/// Converts all of the elements of a `[f32]` slice into a new [`struct@f16`] or [`bf16`] vector.
|
|
///
|
|
/// The conversion operation is vectorized over the slice, meaning the conversion may be more
|
|
/// efficient than converting individual elements on some hardware that supports SIMD
|
|
/// conversions. See [crate documentation][crate] for more information on hardware conversion
|
|
/// support.
|
|
///
|
|
/// # Examples
|
|
/// ```rust
|
|
/// # use half::prelude::*;
|
|
/// let float_values = [1., 2., 3., 4.];
|
|
/// let vec: Vec<f16> = Vec::from_f32_slice(&float_values);
|
|
///
|
|
/// assert_eq!(vec, vec![f16::from_f32(1.), f16::from_f32(2.), f16::from_f32(3.), f16::from_f32(4.)]);
|
|
/// ```
|
|
#[must_use]
|
|
fn from_f32_slice(slice: &[f32]) -> Self;
|
|
|
|
/// Converts all of the elements of a `[f64]` slice into a new [`struct@f16`] or [`bf16`] vector.
|
|
///
|
|
/// The conversion operation is vectorized over the slice, meaning the conversion may be more
|
|
/// efficient than converting individual elements on some hardware that supports SIMD
|
|
/// conversions. See [crate documentation][crate] for more information on hardware conversion
|
|
/// support.
|
|
///
|
|
/// # Examples
|
|
/// ```rust
|
|
/// # use half::prelude::*;
|
|
/// let float_values = [1., 2., 3., 4.];
|
|
/// let vec: Vec<f16> = Vec::from_f64_slice(&float_values);
|
|
///
|
|
/// assert_eq!(vec, vec![f16::from_f64(1.), f16::from_f64(2.), f16::from_f64(3.), f16::from_f64(4.)]);
|
|
/// ```
|
|
#[must_use]
|
|
fn from_f64_slice(slice: &[f64]) -> Self;
|
|
}
|
|
|
|
/// Extensions to [`Vec<u16>`] to support reinterpret operations.
|
|
///
|
|
/// This trait is sealed and cannot be implemented outside of this crate.
|
|
pub trait HalfBitsVecExt: private::SealedHalfBitsVec {
|
|
/// Reinterprets a vector of [`u16`] bits as a vector of [`struct@f16`] or [`bf16`] numbers.
|
|
///
|
|
/// `H` is the type to cast to, and must be either the [`struct@f16`] or [`bf16`] type.
|
|
///
|
|
/// This is a zero-copy operation. The reinterpreted vector has the same memory location as
|
|
/// `self`.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// ```rust
|
|
/// # use half::prelude::*;
|
|
/// let int_buffer = vec![f16::from_f32(1.).to_bits(), f16::from_f32(2.).to_bits(), f16::from_f32(3.).to_bits()];
|
|
/// let float_buffer = int_buffer.reinterpret_into::<f16>();
|
|
///
|
|
/// assert_eq!(float_buffer, [f16::from_f32(1.), f16::from_f32(2.), f16::from_f32(3.)]);
|
|
/// ```
|
|
#[must_use]
|
|
fn reinterpret_into<H>(self) -> Vec<H>
|
|
where
|
|
H: crate::private::SealedHalf;
|
|
}
|
|
|
|
mod private {
|
|
use crate::{bf16, f16};
|
|
#[cfg(feature = "alloc")]
|
|
#[allow(unused_imports)]
|
|
use alloc::vec::Vec;
|
|
|
|
pub trait SealedHalfFloatVec {}
|
|
impl SealedHalfFloatVec for Vec<f16> {}
|
|
impl SealedHalfFloatVec for Vec<bf16> {}
|
|
|
|
pub trait SealedHalfBitsVec {}
|
|
impl SealedHalfBitsVec for Vec<u16> {}
|
|
}
|
|
|
|
impl HalfFloatVecExt for Vec<f16> {
|
|
#[inline]
|
|
fn reinterpret_into(mut self) -> Vec<u16> {
|
|
// An f16 array has same length and capacity as u16 array
|
|
let length = self.len();
|
|
let capacity = self.capacity();
|
|
|
|
// Actually reinterpret the contents of the Vec<f16> as u16,
|
|
// knowing that structs are represented as only their members in memory,
|
|
// which is the u16 part of `f16(u16)`
|
|
let pointer = self.as_mut_ptr() as *mut u16;
|
|
|
|
// Prevent running a destructor on the old Vec<u16>, so the pointer won't be deleted
|
|
mem::forget(self);
|
|
|
|
// Finally construct a new Vec<f16> from the raw pointer
|
|
// SAFETY: We are reconstructing full length and capacity of original vector,
|
|
// using its original pointer, and the size of elements are identical.
|
|
unsafe { Vec::from_raw_parts(pointer, length, capacity) }
|
|
}
|
|
|
|
#[allow(clippy::uninit_vec)]
|
|
fn from_f32_slice(slice: &[f32]) -> Self {
|
|
let mut vec = vec![f16::from_bits(0); slice.len()];
|
|
vec.convert_from_f32_slice(slice);
|
|
vec
|
|
}
|
|
|
|
#[allow(clippy::uninit_vec)]
|
|
fn from_f64_slice(slice: &[f64]) -> Self {
|
|
let mut vec = vec![f16::from_bits(0); slice.len()];
|
|
vec.convert_from_f64_slice(slice);
|
|
vec
|
|
}
|
|
}
|
|
|
|
impl HalfFloatVecExt for Vec<bf16> {
|
|
#[inline]
|
|
fn reinterpret_into(mut self) -> Vec<u16> {
|
|
// An f16 array has same length and capacity as u16 array
|
|
let length = self.len();
|
|
let capacity = self.capacity();
|
|
|
|
// Actually reinterpret the contents of the Vec<f16> as u16,
|
|
// knowing that structs are represented as only their members in memory,
|
|
// which is the u16 part of `f16(u16)`
|
|
let pointer = self.as_mut_ptr() as *mut u16;
|
|
|
|
// Prevent running a destructor on the old Vec<u16>, so the pointer won't be deleted
|
|
mem::forget(self);
|
|
|
|
// Finally construct a new Vec<f16> from the raw pointer
|
|
// SAFETY: We are reconstructing full length and capacity of original vector,
|
|
// using its original pointer, and the size of elements are identical.
|
|
unsafe { Vec::from_raw_parts(pointer, length, capacity) }
|
|
}
|
|
|
|
#[allow(clippy::uninit_vec)]
|
|
fn from_f32_slice(slice: &[f32]) -> Self {
|
|
let mut vec = vec![bf16::from_bits(0); slice.len()];
|
|
vec.convert_from_f32_slice(slice);
|
|
vec
|
|
}
|
|
|
|
#[allow(clippy::uninit_vec)]
|
|
fn from_f64_slice(slice: &[f64]) -> Self {
|
|
let mut vec = vec![bf16::from_bits(0); slice.len()];
|
|
vec.convert_from_f64_slice(slice);
|
|
vec
|
|
}
|
|
}
|
|
|
|
impl HalfBitsVecExt for Vec<u16> {
|
|
// This is safe because all traits are sealed
|
|
#[inline]
|
|
fn reinterpret_into<H>(mut self) -> Vec<H>
|
|
where
|
|
H: crate::private::SealedHalf,
|
|
{
|
|
// An f16 array has same length and capacity as u16 array
|
|
let length = self.len();
|
|
let capacity = self.capacity();
|
|
|
|
// Actually reinterpret the contents of the Vec<u16> as f16,
|
|
// knowing that structs are represented as only their members in memory,
|
|
// which is the u16 part of `f16(u16)`
|
|
let pointer = self.as_mut_ptr() as *mut H;
|
|
|
|
// Prevent running a destructor on the old Vec<u16>, so the pointer won't be deleted
|
|
mem::forget(self);
|
|
|
|
// Finally construct a new Vec<f16> from the raw pointer
|
|
// SAFETY: We are reconstructing full length and capacity of original vector,
|
|
// using its original pointer, and the size of elements are identical.
|
|
unsafe { Vec::from_raw_parts(pointer, length, capacity) }
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod test {
|
|
use super::{HalfBitsVecExt, HalfFloatVecExt};
|
|
use crate::{bf16, f16};
|
|
#[cfg(all(feature = "alloc", not(feature = "std")))]
|
|
use alloc::vec;
|
|
|
|
#[test]
|
|
fn test_vec_conversions_f16() {
|
|
let numbers = vec![f16::E, f16::PI, f16::EPSILON, f16::FRAC_1_SQRT_2];
|
|
let bits = vec![
|
|
f16::E.to_bits(),
|
|
f16::PI.to_bits(),
|
|
f16::EPSILON.to_bits(),
|
|
f16::FRAC_1_SQRT_2.to_bits(),
|
|
];
|
|
let bits_cloned = bits.clone();
|
|
|
|
// Convert from bits to numbers
|
|
let from_bits = bits.reinterpret_into::<f16>();
|
|
assert_eq!(&from_bits[..], &numbers[..]);
|
|
|
|
// Convert from numbers back to bits
|
|
let to_bits = from_bits.reinterpret_into();
|
|
assert_eq!(&to_bits[..], &bits_cloned[..]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_vec_conversions_bf16() {
|
|
let numbers = vec![bf16::E, bf16::PI, bf16::EPSILON, bf16::FRAC_1_SQRT_2];
|
|
let bits = vec![
|
|
bf16::E.to_bits(),
|
|
bf16::PI.to_bits(),
|
|
bf16::EPSILON.to_bits(),
|
|
bf16::FRAC_1_SQRT_2.to_bits(),
|
|
];
|
|
let bits_cloned = bits.clone();
|
|
|
|
// Convert from bits to numbers
|
|
let from_bits = bits.reinterpret_into::<bf16>();
|
|
assert_eq!(&from_bits[..], &numbers[..]);
|
|
|
|
// Convert from numbers back to bits
|
|
let to_bits = from_bits.reinterpret_into();
|
|
assert_eq!(&to_bits[..], &bits_cloned[..]);
|
|
}
|
|
}
|