use bytemuck::Pod;
use crate::byte_storage::CpuStorageBytes;
use fuel_ir::{Error, Result};
pub trait UnaryOp<T: Copy> {
fn apply(x: T) -> T;
}
pub trait UnaryOpCore {
fn f32(x: f32) -> f32;
fn f64(x: f64) -> f64;
fn bf16(x: half::bf16) -> half::bf16 {
half::bf16::from_f32(Self::f32(x.to_f32()))
}
fn f16(x: half::f16) -> half::f16 {
half::f16::from_f32(Self::f32(x.to_f32()))
}
}
impl<O: UnaryOpCore> UnaryOp<f32> for O {
fn apply(x: f32) -> f32 {
<O as UnaryOpCore>::f32(x)
}
}
impl<O: UnaryOpCore> UnaryOp<f64> for O {
fn apply(x: f64) -> f64 {
<O as UnaryOpCore>::f64(x)
}
}
impl<O: UnaryOpCore> UnaryOp<half::bf16> for O {
fn apply(x: half::bf16) -> half::bf16 {
<O as UnaryOpCore>::bf16(x)
}
}
impl<O: UnaryOpCore> UnaryOp<half::f16> for O {
fn apply(x: half::f16) -> half::f16 {
<O as UnaryOpCore>::f16(x)
}
}
pub fn unary<T, U>(name: &str, input: &CpuStorageBytes, output: &mut CpuStorageBytes) -> Result<()>
where
T: Copy + Pod,
U: UnaryOp<T>,
{
let in_bytes = input.len_bytes();
let out_bytes = output.len_bytes();
if in_bytes != out_bytes {
return Err(Error::Msg(format!(
"{name}: input bytes={in_bytes} != output bytes={out_bytes}",
))
.bt());
}
let in_view: &[T] = input.as_slice()?;
let out_view: &mut [T] = output.as_slice_mut()?;
for (i, slot) in out_view.iter_mut().enumerate() {
*slot = U::apply(in_view[i]);
}
Ok(())
}
pub struct Relu;
impl UnaryOpCore for Relu {
fn f32(x: f32) -> f32 {
if x < 0.0 { 0.0 } else { x }
}
fn f64(x: f64) -> f64 {
if x < 0.0 { 0.0 } else { x }
}
fn bf16(x: half::bf16) -> half::bf16 {
if x < half::bf16::ZERO {
half::bf16::ZERO
} else {
x
}
}
fn f16(x: half::f16) -> half::f16 {
if x < half::f16::ZERO {
half::f16::ZERO
} else {
x
}
}
}
pub struct Neg;
impl UnaryOpCore for Neg {
fn f32(x: f32) -> f32 {
-x
}
fn f64(x: f64) -> f64 {
-x
}
fn bf16(x: half::bf16) -> half::bf16 {
-x
}
fn f16(x: half::f16) -> half::f16 {
-x
}
}
pub struct Sqr;
impl UnaryOpCore for Sqr {
fn f32(x: f32) -> f32 {
x * x
}
fn f64(x: f64) -> f64 {
x * x
}
}
pub struct Sqrt;
impl UnaryOpCore for Sqrt {
fn f32(x: f32) -> f32 {
quiet_f32(x.sqrt())
}
fn f64(x: f64) -> f64 {
quiet_f64(x.sqrt())
}
}
pub struct Recip;
impl UnaryOpCore for Recip {
fn f32(x: f32) -> f32 {
1.0 / x
}
fn f64(x: f64) -> f64 {
1.0 / x
}
}
pub struct Abs;
impl UnaryOpCore for Abs {
fn f32(x: f32) -> f32 {
x.abs()
}
fn f64(x: f64) -> f64 {
x.abs()
}
fn bf16(x: half::bf16) -> half::bf16 {
half::bf16::from_bits(x.to_bits() & 0x7FFF)
}
fn f16(x: half::f16) -> half::f16 {
half::f16::from_bits(x.to_bits() & 0x7FFF)
}
}
#[inline]
pub(crate) fn quiet_f32(r: f32) -> f32 {
if r.is_nan() {
f32::from_bits(r.to_bits() | 0x0040_0000)
} else {
r
}
}
#[inline]
pub(crate) fn quiet_f64(r: f64) -> f64 {
if r.is_nan() {
f64::from_bits(r.to_bits() | 0x0008_0000_0000_0000)
} else {
r
}
}
pub struct Tanh;
impl UnaryOpCore for Tanh {
fn f32(x: f32) -> f32 {
quiet_f32(x.tanh())
}
fn f64(x: f64) -> f64 {
quiet_f64(x.tanh())
}
}
pub struct Exp;
impl UnaryOpCore for Exp {
fn f32(x: f32) -> f32 {
quiet_f32(x.exp())
}
fn f64(x: f64) -> f64 {
quiet_f64(x.exp())
}
}
pub struct Log;
impl UnaryOpCore for Log {
fn f32(x: f32) -> f32 {
quiet_f32(x.ln())
}
fn f64(x: f64) -> f64 {
quiet_f64(x.ln())
}
}
pub struct Sin;
impl UnaryOpCore for Sin {
fn f32(x: f32) -> f32 {
quiet_f32(x.sin())
}
fn f64(x: f64) -> f64 {
quiet_f64(x.sin())
}
}
pub struct Cos;
impl UnaryOpCore for Cos {
fn f32(x: f32) -> f32 {
quiet_f32(x.cos())
}
fn f64(x: f64) -> f64 {
quiet_f64(x.cos())
}
}
pub struct Sigmoid;
impl UnaryOpCore for Sigmoid {
fn f32(x: f32) -> f32 {
1.0 / (1.0 + (-x).exp())
}
fn f64(x: f64) -> f64 {
1.0 / (1.0 + (-x).exp())
}
}
pub struct Silu;
impl UnaryOpCore for Silu {
fn f32(x: f32) -> f32 {
x / (1.0 + (-x).exp())
}
fn f64(x: f64) -> f64 {
x / (1.0 + (-x).exp())
}
}
pub struct Step;
impl UnaryOpCore for Step {
fn f32(x: f32) -> f32 {
if x > 0.0 { 1.0 } else { 0.0 }
}
fn f64(x: f64) -> f64 {
if x > 0.0 { 1.0 } else { 0.0 }
}
}
pub struct Floor;
impl UnaryOpCore for Floor {
fn f32(x: f32) -> f32 {
quiet_f32(x.floor())
}
fn f64(x: f64) -> f64 {
quiet_f64(x.floor())
}
}
pub struct Ceil;
impl UnaryOpCore for Ceil {
fn f32(x: f32) -> f32 {
quiet_f32(x.ceil())
}
fn f64(x: f64) -> f64 {
quiet_f64(x.ceil())
}
}
pub struct Round;
impl UnaryOpCore for Round {
fn f32(x: f32) -> f32 {
quiet_f32(x.round_ties_even())
}
fn f64(x: f64) -> f64 {
quiet_f64(x.round_ties_even())
}
}
pub struct Sign;
impl UnaryOpCore for Sign {
fn f32(x: f32) -> f32 {
if x > 0.0 {
1.0
} else if x < 0.0 {
-1.0
} else {
0.0
}
}
fn f64(x: f64) -> f64 {
if x > 0.0 {
1.0
} else if x < 0.0 {
-1.0
} else {
0.0
}
}
}
pub struct Erf;
impl UnaryOpCore for Erf {
fn f32(x: f32) -> f32 {
quiet_f32(libm::erff(x))
}
fn f64(x: f64) -> f64 {
quiet_f64(libm::erf(x))
}
}
pub struct GeluErf;
impl UnaryOpCore for GeluErf {
fn f32(x: f32) -> f32 {
0.5 * x * (1.0 + libm::erff(x * std::f32::consts::FRAC_1_SQRT_2))
}
fn f64(x: f64) -> f64 {
0.5 * x * (1.0 + libm::erf(x * std::f64::consts::FRAC_1_SQRT_2))
}
}
pub struct Rsqrt;
impl UnaryOpCore for Rsqrt {
fn f32(x: f32) -> f32 {
1.0 / x.sqrt()
}
fn f64(x: f64) -> f64 {
1.0 / x.sqrt()
}
}
pub struct GeluTanh;
impl UnaryOpCore for GeluTanh {
fn f32(x: f32) -> f32 {
const COEFF: f32 = 0.797_884_6;
let inner = COEFF * (x + 0.044_715 * x * x * x);
0.5 * x * (1.0 + inner.tanh())
}
fn f64(x: f64) -> f64 {
const COEFF: f64 = 0.797_884_560_802_865_4;
let inner = COEFF * (x + 0.044_715 * x * x * x);
0.5 * x * (1.0 + inner.tanh())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn remaining_libm_ops_return_a_quiet_nan_for_a_signalling_operand() {
let snan32 = f32::from_bits(0x7F80_0001);
let snan64 = f64::from_bits(0x7FF0_0000_0000_0001);
assert!(snan32.is_nan() && snan32.to_bits() & 0x0040_0000 == 0);
assert!(snan64.is_nan() && snan64.to_bits() & 0x0008_0000_0000_0000 == 0);
let q32 = |r: f32| r.is_nan() && r.to_bits() & 0x0040_0000 != 0;
let q64 = |r: f64| r.is_nan() && r.to_bits() & 0x0008_0000_0000_0000 != 0;
assert!(q32(<Erf as UnaryOpCore>::f32(snan32)), "erf f32");
assert!(q32(<Floor as UnaryOpCore>::f32(snan32)), "floor f32");
assert!(q32(<Ceil as UnaryOpCore>::f32(snan32)), "ceil f32");
assert!(q32(<Round as UnaryOpCore>::f32(snan32)), "round f32");
assert!(q32(<Sqrt as UnaryOpCore>::f32(snan32)), "sqrt f32");
assert!(q32(<Rsqrt as UnaryOpCore>::f32(snan32)), "rsqrt f32");
assert!(q32(<Sigmoid as UnaryOpCore>::f32(snan32)), "sigmoid f32");
assert!(q32(<Silu as UnaryOpCore>::f32(snan32)), "silu f32");
assert!(q32(<GeluErf as UnaryOpCore>::f32(snan32)), "gelu_erf f32");
assert!(q32(<GeluTanh as UnaryOpCore>::f32(snan32)), "gelu_tanh f32");
assert!(q64(<Erf as UnaryOpCore>::f64(snan64)), "erf f64");
assert!(q64(<Floor as UnaryOpCore>::f64(snan64)), "floor f64");
assert!(q64(<Ceil as UnaryOpCore>::f64(snan64)), "ceil f64");
assert!(q64(<Round as UnaryOpCore>::f64(snan64)), "round f64");
assert!(q64(<Sqrt as UnaryOpCore>::f64(snan64)), "sqrt f64");
assert!(q64(<Rsqrt as UnaryOpCore>::f64(snan64)), "rsqrt f64");
assert!(q64(<Sigmoid as UnaryOpCore>::f64(snan64)), "sigmoid f64");
assert!(q64(<Silu as UnaryOpCore>::f64(snan64)), "silu f64");
assert!(q64(<GeluErf as UnaryOpCore>::f64(snan64)), "gelu_erf f64");
assert!(q64(<GeluTanh as UnaryOpCore>::f64(snan64)), "gelu_tanh f64");
}
#[test]
fn transcendentals_return_a_quiet_nan_for_a_signalling_operand() {
let snan32 = f32::from_bits(0x7F80_0001);
let snan64 = f64::from_bits(0x7FF0_0000_0000_0001);
assert!(
snan32.is_nan() && snan32.to_bits() & 0x0040_0000 == 0,
"f32 fixture must be signalling"
);
assert!(
snan64.is_nan() && snan64.to_bits() & 0x0008_0000_0000_0000 == 0,
"f64 fixture must be signalling"
);
fn q32(r: f32) -> bool {
r.is_nan() && r.to_bits() & 0x0040_0000 != 0
}
fn q64(r: f64) -> bool {
r.is_nan() && r.to_bits() & 0x0008_0000_0000_0000 != 0
}
assert!(q32(<Exp as UnaryOpCore>::f32(snan32)), "exp f32");
assert!(q32(<Log as UnaryOpCore>::f32(snan32)), "log f32");
assert!(q32(<Sin as UnaryOpCore>::f32(snan32)), "sin f32");
assert!(q32(<Cos as UnaryOpCore>::f32(snan32)), "cos f32");
assert!(q32(<Tanh as UnaryOpCore>::f32(snan32)), "tanh f32");
assert!(q64(<Exp as UnaryOpCore>::f64(snan64)), "exp f64");
assert!(q64(<Log as UnaryOpCore>::f64(snan64)), "log f64");
assert!(q64(<Sin as UnaryOpCore>::f64(snan64)), "sin f64");
assert!(q64(<Cos as UnaryOpCore>::f64(snan64)), "cos f64");
assert!(q64(<Tanh as UnaryOpCore>::f64(snan64)), "tanh f64");
assert_eq!(quiet_f32(1.5), 1.5);
assert_eq!(quiet_f32(f32::INFINITY), f32::INFINITY);
assert_eq!(quiet_f64(-2.5), -2.5);
assert_eq!(
quiet_f32(f32::from_bits(0x7FC0_1234)).to_bits(),
0x7FC0_1234
);
}
#[test]
fn sign_bit_and_select_ops_preserve_a_signalling_nan_in_narrow() {
const QUIET: u16 = 0x0040;
let s = half::bf16::from_bits(0x7F81); assert!(s.is_nan(), "fixture is not a NaN");
assert_eq!(s.to_bits() & QUIET, 0, "fixture must be SIGNALLING");
let promoted = <Sqrt as UnaryOpCore>::bf16(s);
assert_ne!(
promoted.to_bits() & QUIET,
0,
"control failed: the promoting path did NOT quiet, so 'still signalling' below proves nothing -- check the fixture and half's convert"
);
for (name, got) in [
("Neg", <Neg as UnaryOpCore>::bf16(s)),
("Abs", <Abs as UnaryOpCore>::bf16(s)),
("Relu", <Relu as UnaryOpCore>::bf16(s)),
] {
assert!(got.is_nan(), "{name} lost NaN-ness entirely");
assert_eq!(
got.to_bits() & QUIET,
0,
"{name} QUIETED a signalling NaN (0x{:04X} -> 0x{:04X}); it is a sign-bit/select op and must return SOURCE bits, not round an f32 result. See KISS #355 and IEEE 754-2019 §5.5.1.",
s.to_bits(),
got.to_bits()
);
}
}
#[test]
fn narrow_overrides_agree_with_the_promoting_path_on_non_nan_values() {
for bits in [
0x3F80u16, 0xC020, 0x0000, 0x8000, 0x7F7F, 0xFF7F, 0x7F80, 0xFF80,
] {
let x = half::bf16::from_bits(bits);
assert_eq!(
<Neg as UnaryOpCore>::bf16(x).to_bits(),
half::bf16::from_f32(-x.to_f32()).to_bits(),
"Neg diverged from the promoting path on 0x{bits:04X}"
);
assert_eq!(
<Abs as UnaryOpCore>::bf16(x).to_bits(),
half::bf16::from_f32(x.to_f32().abs()).to_bits(),
"Abs diverged from the promoting path on 0x{bits:04X}"
);
assert_eq!(
<Relu as UnaryOpCore>::bf16(x).to_bits(),
half::bf16::from_f32(if x.to_f32() < 0.0 { 0.0 } else { x.to_f32() }).to_bits(),
"Relu diverged from the promoting path on 0x{bits:04X}"
);
}
let nz = half::bf16::from_f32(-0.0);
assert_eq!(
<Relu as UnaryOpCore>::bf16(nz).to_bits(),
nz.to_bits(),
"Relu must PRESERVE -0.0 (pinned 2026-07-14), not normalise it to +0.0"
);
}
#[test]
fn unary_op_neg_f32_negates() {
assert_eq!(<Neg as UnaryOp<f32>>::apply(2.5), -2.5);
assert_eq!(<Neg as UnaryOp<f32>>::apply(-3.0), 3.0);
assert_eq!(<Neg as UnaryOp<f32>>::apply(0.0), 0.0);
}
#[test]
fn unary_op_relu_f32_clamps_negative_to_zero() {
assert_eq!(<Relu as UnaryOp<f32>>::apply(2.5), 2.5);
assert_eq!(<Relu as UnaryOp<f32>>::apply(-3.0), 0.0);
assert_eq!(<Relu as UnaryOp<f32>>::apply(0.0), 0.0);
}
#[test]
fn unary_op_bf16_blanket_routes_through_f32() {
let x = half::bf16::from_f32(1.5);
let got = <Sqr as UnaryOp<half::bf16>>::apply(x).to_f32();
let expect = half::bf16::from_f32(1.5 * 1.5).to_f32();
assert_eq!(got, expect);
}
#[test]
fn unary_op_gelu_tanh_f32_at_zero_is_zero() {
assert!(<GeluTanh as UnaryOp<f32>>::apply(0.0).abs() < 1e-6);
}
#[test]
fn unary_op_gelu_tanh_f32_at_one() {
let got = <GeluTanh as UnaryOp<f32>>::apply(1.0);
assert!((got - 0.841_192).abs() < 1e-3, "got {got}");
}
#[test]
fn unary_op_sign_f32_handles_zero_as_zero() {
assert_eq!(<Sign as UnaryOp<f32>>::apply(0.0), 0.0);
assert_eq!(<Sign as UnaryOp<f32>>::apply(5.0), 1.0);
assert_eq!(<Sign as UnaryOp<f32>>::apply(-5.0), -1.0);
}
#[test]
fn unary_chassis_size_mismatch_errors() {
let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0]);
let mut output = CpuStorageBytes::from_zero_bytes(4); let r = unary::<f32, Neg>("test", &input, &mut output);
assert!(r.is_err());
}
#[test]
fn unary_chassis_walks_all_elements() {
let input = CpuStorageBytes::from_slice(&[1.0_f32, -2.0, 3.0, -4.0]);
let mut output = CpuStorageBytes::from_zero_bytes(input.len_bytes());
unary::<f32, Neg>("test", &input, &mut output).expect("unary neg_f32");
let r: &[f32] = output.as_slice().unwrap();
assert_eq!(r, &[-1.0, 2.0, -3.0, 4.0]);
}
}