use core::fmt;
use core::simd::prelude::*;
use crate::{Uf16, Uf16E5M11, Uf16E6M10, Uf32};
#[cfg(any(target_feature = "avx2", target_feature = "avx512f"))]
pub const UF16_LANES: usize = 8;
#[cfg(not(any(target_feature = "avx2", target_feature = "avx512f")))]
pub const UF16_LANES: usize = 4;
#[cfg(any(target_feature = "avx2", target_feature = "avx512f"))]
pub const UF32_LANES: usize = 4;
#[cfg(not(any(target_feature = "avx2", target_feature = "avx512f")))]
pub const UF32_LANES: usize = 2;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SimdError {
InputLengthMismatch { left: usize, right: usize },
OutputLengthMismatch { input: usize, output: usize },
}
impl fmt::Display for SimdError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::InputLengthMismatch { left, right } => {
write!(formatter, "SIMD input lengths differ ({left} and {right})")
}
Self::OutputLengthMismatch { input, output } => {
write!(
formatter,
"SIMD output length is {output}, expected {input}"
)
}
}
}
}
trait Uf16Layout: Copy {
const EXPONENT_BITS: u32;
const MANTISSA_BITS: u32;
const F32_EXPONENT_BIAS: u32;
fn from_bits(bits: u16) -> Self;
fn to_bits(self) -> u16;
fn from_f32(value: f32) -> Self;
fn to_f32(self) -> f32;
}
impl Uf16Layout for Uf16E5M11 {
const EXPONENT_BITS: u32 = 5;
const MANTISSA_BITS: u32 = 11;
const F32_EXPONENT_BIAS: u32 = 112;
fn from_bits(bits: u16) -> Self {
Self::from_bits(bits)
}
fn to_bits(self) -> u16 {
self.to_bits()
}
fn from_f32(value: f32) -> Self {
Self::from_f32(value)
}
fn to_f32(self) -> f32 {
self.to_f32()
}
}
impl Uf16Layout for Uf16E6M10 {
const EXPONENT_BITS: u32 = 6;
const MANTISSA_BITS: u32 = 10;
const F32_EXPONENT_BIAS: u32 = 96;
fn from_bits(bits: u16) -> Self {
Self::from_bits(bits)
}
fn to_bits(self) -> u16 {
self.to_bits()
}
fn from_f32(value: f32) -> Self {
Self::from_f32(value)
}
fn to_f32(self) -> f32 {
self.to_f32()
}
}
fn output_len(input: usize, output: usize) -> Result<(), SimdError> {
if input == output {
Ok(())
} else {
Err(SimdError::OutputLengthMismatch { input, output })
}
}
fn binary_len(left: usize, right: usize, output: usize) -> Result<(), SimdError> {
if left != right {
return Err(SimdError::InputLengthMismatch { left, right });
}
output_len(left, output)
}
fn uf16_max_exponent<T: Uf16Layout>() -> u32 {
(1 << T::EXPONENT_BITS) - 1
}
fn can_decode_uf16<T: Uf16Layout>(value: T) -> bool {
let exponent = (value.to_bits() as u32 >> T::MANTISSA_BITS) & uf16_max_exponent::<T>();
exponent != 0 && exponent != uf16_max_exponent::<T>()
}
fn can_encode_uf16<T: Uf16Layout>(value: f32) -> bool {
let bits = value.to_bits();
let exponent = (bits >> 23) & 0xff;
let max_normal = uf16_max_exponent::<T>() - 1;
bits >> 31 == 0
&& exponent > T::F32_EXPONENT_BIAS
&& exponent < T::F32_EXPONENT_BIAS + max_normal
}
fn decode_uf16_fast<T: Uf16Layout>(src: &[T]) -> Simd<f32, UF16_LANES> {
debug_assert_eq!(src.len(), UF16_LANES);
debug_assert!(src.iter().copied().all(can_decode_uf16::<T>));
let raw = Simd::<u32, UF16_LANES>::from_array(core::array::from_fn(|lane| {
src[lane].to_bits() as u32
}));
let bits =
(raw << Simd::splat(23 - T::MANTISSA_BITS)) + Simd::splat(T::F32_EXPONENT_BIAS << 23);
Simd::<f32, UF16_LANES>::from_bits(bits)
}
fn encode_uf16_fast<T: Uf16Layout>(src: Simd<f32, UF16_LANES>, dst: &mut [T]) {
debug_assert_eq!(dst.len(), UF16_LANES);
debug_assert!(src.to_array().into_iter().all(can_encode_uf16::<T>));
let bits = src.to_bits();
let fraction = bits & Simd::splat(0x007f_ffff_u32);
let drop = 23 - T::MANTISSA_BITS;
let mantissa = fraction >> Simd::splat(drop);
let discarded = fraction & Simd::splat((1_u32 << drop) - 1);
let rounding =
(discarded + Simd::splat((1_u32 << (drop - 1)) - 1) + (mantissa & Simd::splat(1)))
>> Simd::splat(drop);
let rounded = mantissa + rounding;
let carry = rounded >> Simd::splat(T::MANTISSA_BITS);
let exponent = (bits >> Simd::splat(23)) - Simd::splat(T::F32_EXPONENT_BIAS) + carry;
let raw = (exponent << Simd::splat(T::MANTISSA_BITS))
| (rounded & Simd::splat((1_u32 << T::MANTISSA_BITS) - 1));
for (lane, bits) in raw.to_array().into_iter().enumerate() {
dst[lane] = T::from_bits(bits as u16);
}
}
fn decode_uf16<T: Uf16Layout>(src: &[T], dst: &mut [f32]) -> Result<(), SimdError> {
output_len(src.len(), dst.len())?;
let vector_end = src.len() / UF16_LANES * UF16_LANES;
for offset in (0..vector_end).step_by(UF16_LANES) {
let input = &src[offset..offset + UF16_LANES];
if input.iter().copied().all(can_decode_uf16::<T>) {
decode_uf16_fast(input).copy_to_slice(&mut dst[offset..]);
} else {
for lane in 0..UF16_LANES {
dst[offset + lane] = input[lane].to_f32();
}
}
}
for index in vector_end..src.len() {
dst[index] = src[index].to_f32();
}
Ok(())
}
fn encode_uf16<T: Uf16Layout>(src: &[f32], dst: &mut [T]) -> Result<(), SimdError> {
output_len(src.len(), dst.len())?;
let vector_end = src.len() / UF16_LANES * UF16_LANES;
for offset in (0..vector_end).step_by(UF16_LANES) {
let input = Simd::<f32, UF16_LANES>::from_slice(&src[offset..]);
if input.to_array().into_iter().all(can_encode_uf16::<T>) {
encode_uf16_fast(input, &mut dst[offset..offset + UF16_LANES]);
} else {
for lane in 0..UF16_LANES {
dst[offset + lane] = T::from_f32(src[offset + lane]);
}
}
}
for index in vector_end..src.len() {
dst[index] = T::from_f32(src[index]);
}
Ok(())
}
fn binary_uf16<T: Uf16Layout>(
left: &[T],
right: &[T],
output: &mut [T],
vector: impl Fn(Simd<f32, UF16_LANES>, Simd<f32, UF16_LANES>) -> Simd<f32, UF16_LANES>,
scalar: impl Fn(f32, f32) -> f32,
) -> Result<(), SimdError> {
binary_len(left.len(), right.len(), output.len())?;
let vector_end = left.len() / UF16_LANES * UF16_LANES;
for offset in (0..vector_end).step_by(UF16_LANES) {
let lhs = &left[offset..offset + UF16_LANES];
let rhs = &right[offset..offset + UF16_LANES];
if lhs.iter().copied().all(can_decode_uf16::<T>)
&& rhs.iter().copied().all(can_decode_uf16::<T>)
{
let result = vector(decode_uf16_fast(lhs), decode_uf16_fast(rhs));
if result.to_array().into_iter().all(can_encode_uf16::<T>) {
encode_uf16_fast(result, &mut output[offset..offset + UF16_LANES]);
continue;
}
}
for lane in 0..UF16_LANES {
output[offset + lane] = T::from_f32(scalar(lhs[lane].to_f32(), rhs[lane].to_f32()));
}
}
for index in vector_end..left.len() {
output[index] = T::from_f32(scalar(left[index].to_f32(), right[index].to_f32()));
}
Ok(())
}
pub fn decode_uf16_to_f32(src: &[Uf16], dst: &mut [f32]) -> Result<(), SimdError> {
decode_uf16(src, dst)
}
pub fn encode_f32_to_uf16(src: &[f32], dst: &mut [Uf16]) -> Result<(), SimdError> {
encode_uf16(src, dst)
}
pub fn decode_uf16e6m10_to_f32(src: &[Uf16E6M10], dst: &mut [f32]) -> Result<(), SimdError> {
decode_uf16(src, dst)
}
pub fn encode_f32_to_uf16e6m10(src: &[f32], dst: &mut [Uf16E6M10]) -> Result<(), SimdError> {
encode_uf16(src, dst)
}
pub fn add_uf16(left: &[Uf16], right: &[Uf16], output: &mut [Uf16]) -> Result<(), SimdError> {
binary_uf16(
left,
right,
output,
|left, right| left + right,
|left, right| left + right,
)
}
pub fn sub_uf16(left: &[Uf16], right: &[Uf16], output: &mut [Uf16]) -> Result<(), SimdError> {
binary_uf16(
left,
right,
output,
|left, right| left - right,
|left, right| left - right,
)
}
pub fn mul_uf16(left: &[Uf16], right: &[Uf16], output: &mut [Uf16]) -> Result<(), SimdError> {
binary_uf16(
left,
right,
output,
|left, right| left * right,
|left, right| left * right,
)
}
pub fn div_uf16(left: &[Uf16], right: &[Uf16], output: &mut [Uf16]) -> Result<(), SimdError> {
binary_uf16(
left,
right,
output,
|left, right| left / right,
|left, right| left / right,
)
}
pub fn add_uf16e6m10(
left: &[Uf16E6M10],
right: &[Uf16E6M10],
output: &mut [Uf16E6M10],
) -> Result<(), SimdError> {
binary_uf16(
left,
right,
output,
|left, right| left + right,
|left, right| left + right,
)
}
pub fn sub_uf16e6m10(
left: &[Uf16E6M10],
right: &[Uf16E6M10],
output: &mut [Uf16E6M10],
) -> Result<(), SimdError> {
binary_uf16(
left,
right,
output,
|left, right| left - right,
|left, right| left - right,
)
}
pub fn mul_uf16e6m10(
left: &[Uf16E6M10],
right: &[Uf16E6M10],
output: &mut [Uf16E6M10],
) -> Result<(), SimdError> {
binary_uf16(
left,
right,
output,
|left, right| left * right,
|left, right| left * right,
)
}
pub fn div_uf16e6m10(
left: &[Uf16E6M10],
right: &[Uf16E6M10],
output: &mut [Uf16E6M10],
) -> Result<(), SimdError> {
binary_uf16(
left,
right,
output,
|left, right| left / right,
|left, right| left / right,
)
}
fn can_decode_uf32(value: Uf32) -> bool {
let exponent = value.to_bits() >> 24;
exponent != 0 && exponent != 0xff
}
fn can_encode_uf32(value: f64) -> bool {
let bits = value.to_bits();
let exponent = (bits >> 52) & 0x7ff;
bits >> 63 == 0 && exponent > 896 && exponent < 1150
}
fn decode_uf32_fast(src: &[Uf32]) -> Simd<f64, UF32_LANES> {
debug_assert_eq!(src.len(), UF32_LANES);
debug_assert!(src.iter().copied().all(can_decode_uf32));
let raw = Simd::<u64, UF32_LANES>::from_array(core::array::from_fn(|lane| {
src[lane].to_bits() as u64
}));
let bits = (raw << Simd::splat(28)) + Simd::splat(896_u64 << 52);
Simd::<f64, UF32_LANES>::from_bits(bits)
}
fn encode_uf32_fast(src: Simd<f64, UF32_LANES>, dst: &mut [Uf32]) {
debug_assert_eq!(dst.len(), UF32_LANES);
debug_assert!(src.to_array().into_iter().all(can_encode_uf32));
let bits = src.to_bits();
let fraction = bits & Simd::splat(0x000f_ffff_ffff_ffff_u64);
let mantissa = fraction >> Simd::splat(28);
let discarded = fraction & Simd::splat((1_u64 << 28) - 1);
let rounding = (discarded + Simd::splat((1_u64 << 27) - 1) + (mantissa & Simd::splat(1)))
>> Simd::splat(28);
let rounded = mantissa + rounding;
let carry = rounded >> Simd::splat(24);
let exponent = (bits >> Simd::splat(52)) - Simd::splat(896_u64) + carry;
let raw = (exponent << Simd::splat(24)) | (rounded & Simd::splat(0x00ff_ffff_u64));
for (lane, bits) in raw.to_array().into_iter().enumerate() {
dst[lane] = Uf32::from_bits(bits as u32);
}
}
pub fn decode_uf32_to_f64(src: &[Uf32], dst: &mut [f64]) -> Result<(), SimdError> {
output_len(src.len(), dst.len())?;
let vector_end = src.len() / UF32_LANES * UF32_LANES;
for offset in (0..vector_end).step_by(UF32_LANES) {
let input = &src[offset..offset + UF32_LANES];
if input.iter().copied().all(can_decode_uf32) {
decode_uf32_fast(input).copy_to_slice(&mut dst[offset..]);
} else {
for lane in 0..UF32_LANES {
dst[offset + lane] = input[lane].to_f64();
}
}
}
for index in vector_end..src.len() {
dst[index] = src[index].to_f64();
}
Ok(())
}
pub fn encode_f64_to_uf32(src: &[f64], dst: &mut [Uf32]) -> Result<(), SimdError> {
output_len(src.len(), dst.len())?;
let vector_end = src.len() / UF32_LANES * UF32_LANES;
for offset in (0..vector_end).step_by(UF32_LANES) {
let input = Simd::<f64, UF32_LANES>::from_slice(&src[offset..]);
if input.to_array().into_iter().all(can_encode_uf32) {
encode_uf32_fast(input, &mut dst[offset..offset + UF32_LANES]);
} else {
for lane in 0..UF32_LANES {
dst[offset + lane] = Uf32::from_f64(src[offset + lane]);
}
}
}
for index in vector_end..src.len() {
dst[index] = Uf32::from_f64(src[index]);
}
Ok(())
}
fn binary_uf32(
left: &[Uf32],
right: &[Uf32],
output: &mut [Uf32],
vector: impl Fn(Simd<f64, UF32_LANES>, Simd<f64, UF32_LANES>) -> Simd<f64, UF32_LANES>,
scalar: impl Fn(f64, f64) -> f64,
) -> Result<(), SimdError> {
binary_len(left.len(), right.len(), output.len())?;
let vector_end = left.len() / UF32_LANES * UF32_LANES;
for offset in (0..vector_end).step_by(UF32_LANES) {
let lhs = &left[offset..offset + UF32_LANES];
let rhs = &right[offset..offset + UF32_LANES];
if lhs.iter().copied().all(can_decode_uf32) && rhs.iter().copied().all(can_decode_uf32) {
let result = vector(decode_uf32_fast(lhs), decode_uf32_fast(rhs));
if result.to_array().into_iter().all(can_encode_uf32) {
encode_uf32_fast(result, &mut output[offset..offset + UF32_LANES]);
continue;
}
}
for lane in 0..UF32_LANES {
output[offset + lane] = Uf32::from_f64(scalar(lhs[lane].to_f64(), rhs[lane].to_f64()));
}
}
for index in vector_end..left.len() {
output[index] = Uf32::from_f64(scalar(left[index].to_f64(), right[index].to_f64()));
}
Ok(())
}
pub fn add_uf32(left: &[Uf32], right: &[Uf32], output: &mut [Uf32]) -> Result<(), SimdError> {
binary_uf32(
left,
right,
output,
|left, right| left + right,
|left, right| left + right,
)
}
pub fn sub_uf32(left: &[Uf32], right: &[Uf32], output: &mut [Uf32]) -> Result<(), SimdError> {
binary_uf32(
left,
right,
output,
|left, right| left - right,
|left, right| left - right,
)
}
pub fn mul_uf32(left: &[Uf32], right: &[Uf32], output: &mut [Uf32]) -> Result<(), SimdError> {
binary_uf32(
left,
right,
output,
|left, right| left * right,
|left, right| left * right,
)
}
pub fn div_uf32(left: &[Uf32], right: &[Uf32], output: &mut [Uf32]) -> Result<(), SimdError> {
binary_uf32(
left,
right,
output,
|left, right| left / right,
|left, right| left / right,
)
}
#[cfg(test)]
mod tests {
use super::*;
use std::vec;
use std::vec::Vec;
fn lcg(state: &mut u64) -> u64 {
*state = state
.wrapping_mul(6_364_136_223_846_793_005)
.wrapping_add(1_442_695_040_888_963_407);
*state
}
fn verify_uf16_conversions<T: Uf16Layout>() {
let source: Vec<T> = (u16::MIN..=u16::MAX).map(T::from_bits).collect();
let mut decoded = vec![0.0; source.len()];
decode_uf16(&source, &mut decoded).unwrap();
for (value, actual) in source.iter().copied().zip(decoded) {
assert_eq!(actual.to_bits(), value.to_f32().to_bits());
}
let mut state = 0x6f_1d_5eed_u64;
let mut input = vec![0.0; 32_771];
input[0] = 0.0;
input[1] = -0.0;
input[2] = f32::INFINITY;
input[3] = f32::NEG_INFINITY;
input[4] = f32::NAN;
for value in &mut input[5..] {
*value = f32::from_bits(lcg(&mut state) as u32);
}
let mut encoded = vec![T::from_bits(0); input.len()];
encode_uf16(&input, &mut encoded).unwrap();
for (value, actual) in input.into_iter().zip(encoded) {
assert_eq!(actual.to_bits(), T::from_f32(value).to_bits());
}
}
fn verify_uf16_binary<T: Uf16Layout>() {
let mut state = 0x9a_6d_ef_41_u64;
let left: Vec<T> = (0..(UF16_LANES * 19 + 3))
.map(|_| T::from_bits(lcg(&mut state) as u16))
.collect();
let right: Vec<T> = (0..left.len())
.map(|_| T::from_bits(lcg(&mut state) as u16))
.collect();
let mut output = vec![T::from_bits(0); left.len()];
binary_uf16(
&left,
&right,
&mut output,
|left, right| left + right,
|left, right| left + right,
)
.unwrap();
for ((left, right), actual) in left.iter().zip(&right).zip(&output) {
assert_eq!(
actual.to_bits(),
T::from_f32(left.to_f32() + right.to_f32()).to_bits()
);
}
binary_uf16(
&left,
&right,
&mut output,
|left, right| left - right,
|left, right| left - right,
)
.unwrap();
for ((left, right), actual) in left.iter().zip(&right).zip(&output) {
assert_eq!(
actual.to_bits(),
T::from_f32(left.to_f32() - right.to_f32()).to_bits()
);
}
binary_uf16(
&left,
&right,
&mut output,
|left, right| left * right,
|left, right| left * right,
)
.unwrap();
for ((left, right), actual) in left.iter().zip(&right).zip(&output) {
assert_eq!(
actual.to_bits(),
T::from_f32(left.to_f32() * right.to_f32()).to_bits()
);
}
binary_uf16(
&left,
&right,
&mut output,
|left, right| left / right,
|left, right| left / right,
)
.unwrap();
for ((left, right), actual) in left.iter().zip(&right).zip(&output) {
assert_eq!(
actual.to_bits(),
T::from_f32(left.to_f32() / right.to_f32()).to_bits()
);
}
}
#[test]
fn uf16e5m11_bulk_paths_are_bit_exact() {
verify_uf16_conversions::<Uf16>();
verify_uf16_binary::<Uf16>();
}
#[test]
fn uf16e6m10_bulk_paths_are_bit_exact() {
verify_uf16_conversions::<Uf16E6M10>();
verify_uf16_binary::<Uf16E6M10>();
}
#[test]
fn uf32_bulk_paths_are_bit_exact() {
let mut state = 0x03_2d_99_ef_u64;
let mut source = vec![Uf32::ZERO, Uf32::MIN_POSITIVE, Uf32::INFINITY, Uf32::NAN];
source.extend((0..32_767).map(|_| Uf32::from_bits(lcg(&mut state) as u32)));
let mut decoded = vec![0.0; source.len()];
decode_uf32_to_f64(&source, &mut decoded).unwrap();
for (value, actual) in source.iter().copied().zip(decoded) {
assert_eq!(actual.to_bits(), value.to_f64().to_bits());
}
let mut encoded_input = vec![0.0; 32_771];
encoded_input[0] = 0.0;
encoded_input[1] = -0.0;
encoded_input[2] = f64::INFINITY;
encoded_input[3] = f64::NAN;
for value in &mut encoded_input[4..] {
*value = f64::from_bits(lcg(&mut state));
}
let mut encoded = vec![Uf32::ZERO; encoded_input.len()];
encode_f64_to_uf32(&encoded_input, &mut encoded).unwrap();
for (value, actual) in encoded_input.into_iter().zip(encoded) {
assert_eq!(actual.to_bits(), Uf32::from_f64(value).to_bits());
}
let left: Vec<Uf32> = (0..(UF32_LANES * 19 + 1))
.map(|_| Uf32::from_bits(lcg(&mut state) as u32))
.collect();
let right: Vec<Uf32> = (0..left.len())
.map(|_| Uf32::from_bits(lcg(&mut state) as u32))
.collect();
let mut output = vec![Uf32::ZERO; left.len()];
macro_rules! assert_uf32_binary {
($vector:expr, $scalar:expr) => {{
binary_uf32(&left, &right, &mut output, $vector, $scalar).unwrap();
for ((left, right), actual) in left.iter().zip(&right).zip(&output) {
assert_eq!(
actual.to_bits(),
Uf32::from_f64($scalar(left.to_f64(), right.to_f64())).to_bits()
);
}
}};
}
assert_uf32_binary!(|left, right| left + right, |left: f64, right: f64| left
+ right);
assert_uf32_binary!(|left, right| left - right, |left: f64, right: f64| left
- right);
assert_uf32_binary!(|left, right| left * right, |left: f64, right: f64| left
* right);
assert_uf32_binary!(|left, right| left / right, |left: f64, right: f64| left
/ right);
}
#[test]
fn bulk_operations_reject_mismatched_planes() {
let input = [Uf16::ONE; 2];
let other = [Uf16::ONE; 1];
let mut output = [Uf16::ZERO; 2];
assert_eq!(
add_uf16(&input, &other, &mut output),
Err(SimdError::InputLengthMismatch { left: 2, right: 1 })
);
let mut short = [0.0; 1];
assert_eq!(
decode_uf16_to_f32(&input, &mut short),
Err(SimdError::OutputLengthMismatch {
input: 2,
output: 1
})
);
}
}