use super::{FixedMatrix, FixedVector};
use super::kernels;
use crate::fixed_point::fixed_num::{Compute, Fixed, FixedNum};
use crate::fixed_point::width::{Storage, WideInt};
use crate::fixed_point::core_types::errors::OverflowDetected;
#[inline]
fn sum_add<S: Storage, const F: u32>(a: Compute<Fixed<S, F>>, b: Compute<Fixed<S, F>>) -> Compute<Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
Fixed::<S, F>::compute_checked_add(a, b).expect("fused: sum exceeds the compute tier")
}
#[inline]
fn compute_zero<S: Storage, const F: u32>() -> Compute<Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
Fixed::<S, F>::upscale_to_compute(Fixed::ZERO.raw())
}
#[inline]
fn compute_one<S: Storage, const F: u32>() -> Compute<Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
Fixed::<S, F>::upscale_to_compute(Fixed::one().raw())
}
pub fn sqrt_sum_sq<S: Storage, const F: u32>(values: &[Fixed<S, F>]) -> Fixed<S, F>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let mut acc = compute_zero::<S, F>();
for v in values {
let vc = Fixed::<S, F>::upscale_to_compute(v.raw());
acc = sum_add(acc, Fixed::<S, F>::compute_multiply(vc, vc));
}
Fixed::from_raw(kernels::round_to_storage::<Fixed<S, F>>(Fixed::<S, F>::sqrt_at_compute_tier(acc)))
}
pub fn inv_sqrt_sum_sq<S: Storage, const F: u32>(values: &[Fixed<S, F>]) -> Fixed<S, F>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let mut acc = compute_zero::<S, F>();
for v in values {
let vc = Fixed::<S, F>::upscale_to_compute(v.raw());
acc = sum_add(acc, Fixed::<S, F>::compute_multiply(vc, vc));
}
let s = Fixed::<S, F>::sqrt_at_compute_tier(acc);
let inv = Fixed::<S, F>::compute_divide(Fixed::<S, F>::make_compute_int(1), s)
.expect("inv_sqrt_sum_sq: zero norm has no reciprocal");
Fixed::from_raw(kernels::round_to_storage::<Fixed<S, F>>(inv))
}
pub fn euclidean_distance<S: Storage, const F: u32>(a: &[Fixed<S, F>], b: &[Fixed<S, F>]) -> Fixed<S, F>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
assert_eq!(a.len(), b.len(), "euclidean_distance: dimension mismatch");
let mut acc = compute_zero::<S, F>();
for i in 0..a.len() {
let da = Fixed::<S, F>::upscale_to_compute(a[i].raw());
let db = Fixed::<S, F>::upscale_to_compute(b[i].raw());
let diff = Fixed::<S, F>::compute_subtract(da, db);
acc = sum_add(acc, Fixed::<S, F>::compute_multiply(diff, diff));
}
Fixed::from_raw(kernels::round_to_storage::<Fixed<S, F>>(Fixed::<S, F>::sqrt_at_compute_tier(acc)))
}
pub fn euclidean_distance_squared<S: Storage, const F: u32>(a: &[Fixed<S, F>], b: &[Fixed<S, F>]) -> Fixed<S, F>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
assert_eq!(a.len(), b.len(), "euclidean_distance_squared: dimension mismatch");
let mut acc = compute_zero::<S, F>();
for i in 0..a.len() {
let da = Fixed::<S, F>::upscale_to_compute(a[i].raw());
let db = Fixed::<S, F>::upscale_to_compute(b[i].raw());
let diff = Fixed::<S, F>::compute_subtract(da, db);
acc = sum_add(acc, Fixed::<S, F>::compute_multiply(diff, diff));
}
Fixed::from_raw(kernels::round_to_storage::<Fixed<S, F>>(acc))
}
pub fn dot<S: Storage, const F: u32>(a: &[Fixed<S, F>], b: &[Fixed<S, F>]) -> Fixed<S, F>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
assert_eq!(a.len(), b.len(), "dot: dimension mismatch");
let mut acc = compute_zero::<S, F>();
for i in 0..a.len() {
let da = Fixed::<S, F>::upscale_to_compute(a[i].raw());
let db = Fixed::<S, F>::upscale_to_compute(b[i].raw());
acc = sum_add(acc, Fixed::<S, F>::compute_multiply(da, db));
}
Fixed::from_raw(kernels::round_to_storage::<Fixed<S, F>>(acc))
}
pub fn quadratic_form<S: Storage, const F: u32>(v: &FixedVector<Fixed<S, F>>, m: &FixedMatrix<Fixed<S, F>>) -> Fixed<S, F>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
try_quadratic_form(v, m).expect("quadratic_form: storage overflow")
}
pub fn try_quadratic_form<S: Storage, const F: u32>(v: &FixedVector<Fixed<S, F>>, m: &FixedMatrix<Fixed<S, F>>) -> Result<Fixed<S, F>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
Ok(Fixed::from_raw(kernels::narrow_triple_nearest::<Fixed<S, F>>(kernels::quadratic_form_exact(Fixed::raw_slice(v.as_slice()), m)?)?))
}
pub fn mobius_denominator_sq<S: Storage, const F: u32>(p: &[Fixed<S, F>], q: &[Fixed<S, F>]) -> Fixed<S, F>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
assert_eq!(p.len(), q.len(), "mobius_denominator_sq: dimension mismatch");
let mut dot_acc = compute_zero::<S, F>();
let mut p_sq = compute_zero::<S, F>();
let mut q_sq = compute_zero::<S, F>();
for i in 0..p.len() {
let dp = Fixed::<S, F>::upscale_to_compute(p[i].raw());
let dq = Fixed::<S, F>::upscale_to_compute(q[i].raw());
dot_acc = sum_add(dot_acc, Fixed::<S, F>::compute_multiply(dp, dq));
p_sq = sum_add(p_sq, Fixed::<S, F>::compute_multiply(dp, dp));
q_sq = sum_add(q_sq, Fixed::<S, F>::compute_multiply(dq, dq));
}
let one = compute_one::<S, F>();
let two_dot = sum_add(dot_acc, dot_acc);
let result = sum_add(
Fixed::<S, F>::compute_subtract(one, two_dot),
Fixed::<S, F>::compute_multiply(p_sq, q_sq),
);
Fixed::from_raw(kernels::round_to_storage::<Fixed<S, F>>(result))
}
pub fn softmax<S: Storage, const F: u32>(scores: &[Fixed<S, F>]) -> Result<Vec<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
if scores.is_empty() {
return Ok(vec![]);
}
let mut max_raw = scores[0].raw();
for s in &scores[1..] {
if s.raw() > max_raw {
max_raw = s.raw();
}
}
let max_compute = Fixed::<S, F>::upscale_to_compute(max_raw);
let mut exp_values: Vec<Compute<Fixed<S, F>>> = Vec::with_capacity(scores.len());
let mut sum = compute_zero::<S, F>();
for s in scores {
let s_compute = Fixed::<S, F>::upscale_to_compute(s.raw());
let shifted = Fixed::<S, F>::compute_subtract(s_compute, max_compute);
let e = Fixed::<S, F>::exp_at_compute_tier(shifted);
sum = Fixed::<S, F>::compute_checked_add(sum, e)?;
exp_values.push(e);
}
if Fixed::<S, F>::compute_is_zero(&sum) {
return Err(OverflowDetected::DivisionByZero);
}
let mut result = Vec::with_capacity(scores.len());
for e in &exp_values {
let normalized = Fixed::<S, F>::compute_divide(*e, sum)?;
result.push(Fixed::from_raw(Fixed::<S, F>::downscale_to_storage(normalized)?));
}
Ok(result)
}
pub fn rms_norm_factor<S: Storage, const F: u32>(values: &[Fixed<S, F>], eps: Fixed<S, F>) -> Result<Fixed<S, F>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
rms_norm_factor_at_compute(values, Fixed::<S, F>::upscale_to_compute(eps.raw()))
}
pub fn rms_norm_factor_eps_wide<S: Storage, const F: u32>(values: &[Fixed<S, F>], eps_q64: i128) -> Result<Fixed<S, F>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
rms_norm_factor_at_compute(values, q64_to_compute(eps_q64)?)
}
fn q64_to_compute<S: Storage, const F: u32>(x: i128) -> Result<Compute<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let cf = 2 * F;
if cf < 64 {
let shift = 64 - cf;
let rounded = (x >> shift) + ((x >> (shift - 1)) & 1);
Compute::<Fixed<S, F>>::from_i128_checked(rounded).ok_or(OverflowDetected::TierOverflow)
} else {
Ok(Compute::<Fixed<S, F>>::from_i128(x).shl_bits(cf - 64))
}
}
fn rms_norm_factor_at_compute<S: Storage, const F: u32>(values: &[Fixed<S, F>], eps_compute: Compute<Fixed<S, F>>) -> Result<Fixed<S, F>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
if values.is_empty() {
return Err(OverflowDetected::DivisionByZero);
}
let mut sum_sq = compute_zero::<S, F>();
for v in values {
let vc = Fixed::<S, F>::upscale_to_compute(v.raw());
sum_sq = Fixed::<S, F>::compute_checked_add(sum_sq, Fixed::<S, F>::compute_multiply(vc, vc))?;
}
let mean = Fixed::<S, F>::compute_div_count(sum_sq, values.len())?;
let mean_eps = Fixed::<S, F>::compute_checked_add(mean, eps_compute)?;
if Fixed::<S, F>::compute_is_negative(&mean_eps) {
return Err(OverflowDetected::DomainError);
}
let root = Fixed::<S, F>::sqrt_at_compute_tier(mean_eps);
if Fixed::<S, F>::compute_is_zero(&root) {
return Err(OverflowDetected::DivisionByZero);
}
let inv = Fixed::<S, F>::compute_divide(compute_one::<S, F>(), root)?;
Ok(Fixed::from_raw(Fixed::<S, F>::downscale_to_storage(inv)?))
}
pub fn silu<S: Storage, const F: u32>(x: Fixed<S, F>) -> Fixed<S, F>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let x_compute = Fixed::<S, F>::upscale_to_compute(x.raw());
let neg_x = Fixed::<S, F>::compute_negate(x_compute);
let exp_neg = Fixed::<S, F>::exp_at_compute_tier(neg_x);
let one_plus_exp = match Fixed::<S, F>::compute_checked_add(compute_one::<S, F>(), exp_neg) {
Ok(v) => v,
Err(_) => return Fixed::ZERO,
};
if Fixed::<S, F>::compute_is_zero(&one_plus_exp) {
return Fixed::ZERO;
}
match Fixed::<S, F>::compute_divide(x_compute, one_plus_exp) {
Ok(result) => Fixed::from_raw(kernels::round_to_storage::<Fixed<S, F>>(result)),
Err(_) => Fixed::ZERO,
}
}
pub fn softmax_mix<S: Storage, const F: u32>(
scores: &[Fixed<S, F>],
values: &[&[Fixed<S, F>]],
) -> Result<(Vec<Fixed<S, F>>, Vec<Fixed<S, F>>), OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
assert_eq!(scores.len(), values.len(), "softmax_mix: scores/values length mismatch");
let dim = values.first().map_or(0, |v| v.len());
softmax_mix_core(scores, dim, true, |j| {
let v = values[j];
assert_eq!(v.len(), dim, "softmax_mix: value row {j} has length {}, expected {dim}", v.len());
v
})
}
pub fn softmax_mix_values<S: Storage, const F: u32>(scores: &[Fixed<S, F>], values: &[&[Fixed<S, F>]]) -> Result<Vec<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
assert_eq!(scores.len(), values.len(), "softmax_mix: scores/values length mismatch");
let dim = values.first().map_or(0, |v| v.len());
softmax_mix_core(scores, dim, false, |j| {
let v = values[j];
assert_eq!(v.len(), dim, "softmax_mix: value row {j} has length {}, expected {dim}", v.len());
v
})
.map(|(out, _)| out)
}
pub fn softmax_mix_flat<S: Storage, const F: u32>(
scores: &[Fixed<S, F>],
values_flat: &[Fixed<S, F>],
dim: usize,
) -> Result<(Vec<Fixed<S, F>>, Vec<Fixed<S, F>>), OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
assert_eq!(values_flat.len(), scores.len() * dim, "softmax_mix_flat: values_flat length is not scores.len() * dim");
softmax_mix_core(scores, dim, true, |j| &values_flat[j * dim..(j + 1) * dim])
}
pub fn softmax_mix_flat_into<S: Storage, const F: u32>(
scores: &[Fixed<S, F>],
values_flat: &[Fixed<S, F>],
dim: usize,
out: &mut [Fixed<S, F>],
weights: &mut [Fixed<S, F>],
) -> Result<(), OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
assert_eq!(values_flat.len(), scores.len() * dim, "softmax_mix_flat_into: values_flat length is not scores.len() * dim");
assert_eq!(out.len(), dim, "softmax_mix_flat_into: out length is not dim");
assert_eq!(weights.len(), scores.len(), "softmax_mix_flat_into: weights length is not scores.len()");
if scores.is_empty() {
return Ok(());
}
softmax_mix_core_into(scores, dim, out, Some(weights), |j| &values_flat[j * dim..(j + 1) * dim])
}
pub fn softmax_mix_flat_values_into<S: Storage, const F: u32>(
scores: &[Fixed<S, F>],
values_flat: &[Fixed<S, F>],
dim: usize,
out: &mut [Fixed<S, F>],
) -> Result<(), OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
assert_eq!(values_flat.len(), scores.len() * dim, "softmax_mix_flat_values_into: values_flat length is not scores.len() * dim");
assert_eq!(out.len(), dim, "softmax_mix_flat_values_into: out length is not dim");
if scores.is_empty() {
return Ok(());
}
softmax_mix_core_into(scores, dim, out, None, |j| &values_flat[j * dim..(j + 1) * dim])
}
pub fn softmax_mix_flat_values<S: Storage, const F: u32>(scores: &[Fixed<S, F>], values_flat: &[Fixed<S, F>], dim: usize) -> Result<Vec<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
assert_eq!(values_flat.len(), scores.len() * dim, "softmax_mix_flat: values_flat length is not scores.len() * dim");
softmax_mix_core(scores, dim, false, |j| &values_flat[j * dim..(j + 1) * dim]).map(|(out, _)| out)
}
fn softmax_mix_core<'a, S: Storage, const F: u32>(
scores: &[Fixed<S, F>],
dim: usize,
want_weights: bool,
row: impl Fn(usize) -> &'a [Fixed<S, F>],
) -> Result<(Vec<Fixed<S, F>>, Vec<Fixed<S, F>>), OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
if scores.is_empty() {
return Ok((vec![], vec![]));
}
let mut out = vec![Fixed::ZERO; dim];
let mut weights = vec![Fixed::ZERO; if want_weights { scores.len() } else { 0 }];
softmax_mix_core_into(scores, dim, &mut out, want_weights.then_some(&mut weights[..]), row)?;
Ok((out, weights))
}
fn softmax_mix_core_into<'a, S: Storage, const F: u32>(
scores: &[Fixed<S, F>],
dim: usize,
out: &mut [Fixed<S, F>],
weights: Option<&mut [Fixed<S, F>]>,
row: impl Fn(usize) -> &'a [Fixed<S, F>],
) -> Result<(), OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let mut max_raw = scores[0].raw();
for s in &scores[1..] {
if s.raw() > max_raw {
max_raw = s.raw();
}
}
let max_compute = Fixed::<S, F>::upscale_to_compute(max_raw);
let mut exp_values: Vec<Compute<Fixed<S, F>>> = Vec::with_capacity(scores.len());
let mut sum = compute_zero::<S, F>();
for s in scores {
let shifted = Fixed::<S, F>::compute_subtract(Fixed::<S, F>::upscale_to_compute(s.raw()), max_compute);
let e = Fixed::<S, F>::exp_at_compute_tier(shifted);
sum = Fixed::<S, F>::compute_checked_add(sum, e)?;
exp_values.push(e);
}
if Fixed::<S, F>::compute_is_zero(&sum) {
return Err(OverflowDetected::DivisionByZero);
}
let mut num: Vec<Compute<Fixed<S, F>>> = vec![compute_zero::<S, F>(); dim];
let fast = S::mix_numerators_fast::<F>(&exp_values, &|j| Fixed::raw_slice(row(j)), &mut num);
if !fast {
for (j, &e) in exp_values.iter().enumerate() {
let v = row(j);
for d in 0..dim {
num[d] = Fixed::<S, F>::compute_checked_add(num[d], Fixed::<S, F>::compute_multiply(e, Fixed::<S, F>::upscale_to_compute(v[d].raw())))?;
}
}
}
for (o, n) in out.iter_mut().zip(&num) {
*o = Fixed::from_raw(Fixed::<S, F>::downscale_to_storage(Fixed::<S, F>::compute_divide(*n, sum)?)?);
}
if let Some(weights) = weights {
for (w, e) in weights.iter_mut().zip(&exp_values) {
*w = Fixed::from_raw(Fixed::<S, F>::downscale_to_storage(Fixed::<S, F>::compute_divide(*e, sum)?)?);
}
}
Ok(())
}
#[inline]
fn gate_sigmoid<S: Storage, const F: u32>(gate: Fixed<S, F>) -> Compute<Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let g = Fixed::<S, F>::upscale_to_compute(gate.raw());
let one = compute_one::<S, F>();
let (num, e) = if Fixed::<S, F>::compute_is_negative(&g) {
let e = Fixed::<S, F>::exp_at_compute_tier(g);
(e, e)
} else {
(one, Fixed::<S, F>::exp_at_compute_tier(Fixed::<S, F>::compute_negate(g)))
};
Fixed::<S, F>::compute_divide(num, sum_add(one, e)).expect("sigmoid: denominator in [1, 2]")
}
pub fn sigmoid_mul<S: Storage, const F: u32>(x: Fixed<S, F>, gate: Fixed<S, F>) -> Fixed<S, F>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
if let Some(r) = S::sigmoid_mul_fast::<F>(x.raw(), gate.raw()) {
return Fixed::from_raw(r);
}
let p = S::widen_product(S::widen(x.raw()), gate_sigmoid(gate));
Fixed::from_raw(kernels::narrow_triple_nearest::<Fixed<S, F>>(p).expect("sigmoid_mul: |x * sigmoid| <= |x|"))
}
pub fn sigmoid_mul_slice<S: Storage, const F: u32>(x: &[Fixed<S, F>], gate: &[Fixed<S, F>]) -> Vec<Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
assert_eq!(x.len(), gate.len(), "sigmoid_mul_slice: length mismatch");
x.iter().zip(gate).map(|(&x, &g)| sigmoid_mul(x, g)).collect()
}
pub fn sigmoid_mul_in_place<S: Storage, const F: u32>(x: &mut [Fixed<S, F>], gate: &[Fixed<S, F>])
where
Fixed<S, F>: FixedNum<Storage = S>,
{
assert_eq!(x.len(), gate.len(), "sigmoid_mul_in_place: length mismatch");
for (x, &g) in x.iter_mut().zip(gate) {
*x = sigmoid_mul(*x, g);
}
}
pub fn entropy<S: Storage, const F: u32>(weights: &[Fixed<S, F>]) -> Result<Fixed<S, F>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
if weights.iter().any(|w| w.is_negative()) {
return Err(OverflowDetected::DomainError);
}
if let Some(r) = S::entropy_fast::<F>(Fixed::raw_slice(weights)) {
return r.map(Fixed::<S, F>::from_raw);
}
let mut sum = compute_zero::<S, F>();
for w in weights {
if w.is_zero() {
continue;
}
let wc = Fixed::<S, F>::upscale_to_compute(w.raw());
sum = Fixed::<S, F>::compute_checked_add(sum, Fixed::<S, F>::compute_checked_multiply(wc, Fixed::<S, F>::ln_at_compute_tier(wc))?)?;
}
Ok(Fixed::from_raw(Fixed::<S, F>::downscale_to_storage(Fixed::<S, F>::compute_checked_negate(sum)?)?))
}
pub fn dot_many<S: Storage, const F: u32>(query: &[Fixed<S, F>], keys_flat: &[Fixed<S, F>], dim: usize) -> Vec<Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
assert_eq!(query.len(), dim, "dot_many: query length is not dim");
if dim == 0 {
assert!(keys_flat.is_empty(), "dot_many: keys_flat must be empty when dim is 0");
return Vec::new();
}
assert_eq!(keys_flat.len() % dim, 0, "dot_many: keys_flat length is not a multiple of dim");
let mut out = vec![Fixed::ZERO; keys_flat.len() / dim];
dot_many_into(query, keys_flat, dim, &mut out);
out
}
pub fn dot_many_into<S: Storage, const F: u32>(query: &[Fixed<S, F>], keys_flat: &[Fixed<S, F>], dim: usize, out: &mut [Fixed<S, F>])
where
Fixed<S, F>: FixedNum<Storage = S>,
{
assert_eq!(query.len(), dim, "dot_many_into: query length is not dim");
assert_eq!(keys_flat.len(), out.len() * dim, "dot_many_into: keys_flat length is not out.len() * dim");
if dim == 0 {
return;
}
if S::dot_many_into_fast::<F>(Fixed::raw_slice(query), Fixed::raw_slice(keys_flat), dim, Fixed::raw_slice_mut(out)) {
return;
}
for (o, key) in out.iter_mut().zip(keys_flat.chunks_exact(dim)) {
*o = kernels::dot(query, key);
}
}
pub fn rms_norm<S: Storage, const F: u32>(x: &[Fixed<S, F>], weight: &[Fixed<S, F>], eps_q64: i128) -> Result<Vec<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
assert_eq!(x.len(), weight.len(), "rms_norm: x/weight length mismatch");
let (inv, shift) = rms_reciprocal(x, eps_q64)?;
let mut out = x.to_vec();
rms_apply_in_place(&mut out, weight, inv, shift)?;
Ok(out)
}
pub fn rms_norm_in_place<S: Storage, const F: u32>(x: &mut [Fixed<S, F>], weight: &[Fixed<S, F>], eps_q64: i128) -> Result<(), OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
assert_eq!(x.len(), weight.len(), "rms_norm: x/weight length mismatch");
let (inv, shift) = rms_reciprocal(x, eps_q64)?;
rms_apply_in_place(x, weight, inv, shift)
}
fn rms_apply_in_place<S: Storage, const F: u32>(x: &mut [Fixed<S, F>], w: &[Fixed<S, F>], inv: Compute<Fixed<S, F>>, shift: u32) -> Result<(), OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
if let Some(r) = S::rms_apply_in_place_fast::<F>(Fixed::raw_slice_mut(x), Fixed::raw_slice(w), inv, shift) {
return r;
}
for (v, &w) in x.iter_mut().zip(w) {
let product = S::widen_product(S::mul_wide(v.raw(), w.raw()), inv);
*v = Fixed::from_raw(kernels::narrow_shifted_nearest::<Fixed<S, F>>(product, shift)?);
}
Ok(())
}
fn rms_reciprocal<S: Storage, const F: u32>(values: &[Fixed<S, F>], eps_q64: i128) -> Result<(Compute<Fixed<S, F>>, u32), OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
if values.is_empty() {
return Err(OverflowDetected::DivisionByZero);
}
if let Some(r) = S::rms_reciprocal_fast::<F>(Fixed::raw_slice(values), eps_q64) {
return r;
}
let f = F;
let n = i64::try_from(values.len()).map_err(|_| OverflowDetected::TierOverflow)?;
let mut total = Fixed::<S, F>::compute_checked_multiply(q64_to_compute(eps_q64)?, Fixed::<S, F>::make_compute_int(n))?;
for v in values {
total = Fixed::<S, F>::compute_checked_add(total, S::mul_wide(v.raw(), v.raw()))?;
}
if Fixed::<S, F>::compute_is_negative(&total) {
return Err(OverflowDetected::DomainError);
}
if Fixed::<S, F>::compute_is_zero(&total) {
return Err(OverflowDetected::DivisionByZero);
}
let up = (4 * f - 1 - kernels::compute_bit_length::<Fixed<S, F>>(total)) & !1;
let mean = Fixed::<S, F>::compute_div_count(total.shl_bits(up), values.len())?;
let excess = kernels::compute_bit_length::<Fixed<S, F>>(mean) as i64 - (2 * f as i64 + 1);
let (normalised, exponent) = if excess >= 0 {
let down = (excess & !1) as u32;
(mean.shr_bits(down), up as i64 - down as i64)
} else {
let more = ((1 - excess) & !1) as u32;
(mean.shl_bits(more), up as i64 + more as i64)
};
let root = Fixed::<S, F>::sqrt_at_compute_tier(normalised);
let inv = Fixed::<S, F>::compute_divide(compute_one::<S, F>(), root)?; let shift = 3 * f as i64 - exponent / 2;
u32::try_from(shift).map(|shift| (inv, shift)).map_err(|_| OverflowDetected::TierOverflow)
}
#[cfg(test)]
mod tests {
macro_rules! suite {
($m:ident, $S:ty, $F:literal, $W:ident) => {
mod $m {
#![allow(dead_code, unused_imports, unreachable_code)]
use super::super::*;
type FixedPoint = crate::fixed_point::Fixed<$S, $F>;
type BinaryStorage = $S;
type ComputeStorage = <$S as crate::fixed_point::Storage>::Compute;
const FRAC_BITS: u32 = $F;
const BITS: u32 = <$S as crate::fixed_point::WideInt>::BITS;
use super::super::*;
fn fp(s: &str) -> FixedPoint {
if s.starts_with('-') { -FixedPoint::from_str(&s[1..]) }
else { FixedPoint::from_str(s) }
}
fn at_least_one_unit(t: FixedPoint) -> FixedPoint {
let mut u = FixedPoint::one();
for _ in 0..FRAC_BITS { u = u / FixedPoint::from_int(2); }
if u > t { u } else { t }
}
fn loose() -> FixedPoint { at_least_one_unit(fp("0.001")) }
fn tight() -> FixedPoint {
if BITS == 32 { at_least_one_unit(fp("0.001")) } else { fp("0.000000001") }
}
#[test]
fn test_sqrt_sum_sq_basic() {
let vals = [fp("3"), fp("4")];
let result = sqrt_sum_sq(&vals);
let diff = (result - fp("5")).abs();
assert!(diff < tight(), "sqrt(3²+4²) = {}, expected 5", result);
}
#[test]
fn test_sqrt_sum_sq_single() {
let vals = [fp("7")];
let result = sqrt_sum_sq(&vals);
let diff = (result - fp("7")).abs();
assert!(diff < tight(), "sqrt(7²) = {}, expected 7", result);
}
#[test]
fn test_euclidean_distance_basic() {
let a = [FixedPoint::ZERO, FixedPoint::ZERO];
let b = [fp("3"), fp("4")];
let dist = euclidean_distance(&a, &b);
let diff = (dist - fp("5")).abs();
assert!(diff < tight(), "dist([0,0],[3,4]) = {}, expected 5", dist);
}
#[test]
fn test_euclidean_distance_same_point() {
let a = [fp("1"), fp("2"), fp("3")];
let dist = euclidean_distance(&a, &a);
assert!(dist.is_zero() || dist.abs() < tight(),
"distance to self should be 0, got {}", dist);
}
#[test]
fn test_euclidean_distance_squared_matches_distance() {
let cases: [(&[FixedPoint; 2], &[FixedPoint; 2]); 3] = [
(&[FixedPoint::ZERO, FixedPoint::ZERO], &[fp("3"), fp("4")]),
(&[fp("0.25"), fp("-0.5")], &[fp("-0.125"), fp("0.75")]),
(&[fp("3"), fp("-4")], &[fp("-3"), fp("4")]),
];
for (a, b) in cases {
let sq = euclidean_distance_squared(a, b);
let d = euclidean_distance(a, b);
let diff = (sq - d * d).abs();
let tol = (d.abs() + fp("1")) * tight();
assert!(diff < tol,
"dist_sq {} vs dist² {} diverged", sq, d * d);
}
let same = [fp("1"), fp("2")];
assert!(euclidean_distance_squared(&same, &same).abs() < tight());
}
on_words!($W in [w32] {
#[test]
#[should_panic(expected = "exceeds storage tier")]
fn test_euclidean_distance_squared_overflow_panics() {
let a = [FixedPoint::from_raw(i32::MAX)];
let b = [FixedPoint::ZERO];
let _ = euclidean_distance_squared(&a, &b);
}
});
#[test]
fn test_dot_basic() {
let a = [fp("1"), fp("2"), fp("3")];
let b = [fp("4"), fp("5"), fp("6")];
assert!((dot(&a, &b) - fp("32")).abs() < tight());
let e1 = [fp("1"), FixedPoint::ZERO];
let e2 = [FixedPoint::ZERO, fp("1")];
assert!(dot(&e1, &e2).abs() < tight());
let c = [fp("-0.5"), fp("0.25")];
let d = [fp("0.5"), fp("0.25")];
assert!((dot(&c, &d) - fp("-0.1875")).abs() < tight());
}
#[test]
fn test_mobius_denominator_sq() {
let p = [fp("0.3"), fp("0.4")];
let q = [fp("-0.2"), fp("0.5")];
let quant_tol = tight() + tight();
let den = mobius_denominator_sq(&p, &q);
assert!((den - fp("0.7925")).abs() < quant_tol,
"mobius_denominator_sq = {}, expected 0.7925", den);
let o = [FixedPoint::ZERO, FixedPoint::ZERO];
assert!((mobius_denominator_sq(&o, &o) - fp("1")).abs() < tight());
let den_pp = mobius_denominator_sq(&p, &p);
let w = fp("1") - fp("0.25");
assert!((den_pp - w * w).abs() < quant_tol);
}
#[test]
fn test_softmax_uniform() {
let scores = vec![fp("1"); 4];
let result = softmax(&scores).unwrap();
let expected = fp("0.25");
for (i, w) in result.iter().enumerate() {
let diff = (*w - expected).abs();
assert!(diff < loose(), "softmax[{}] = {}, expected 0.25", i, w);
}
}
#[test]
fn test_softmax_sums_to_one() {
let scores = vec![fp("1"), fp("2"), fp("3"), fp("4")];
let result = softmax(&scores).unwrap();
let sum: FixedPoint = result.iter().copied().fold(FixedPoint::ZERO, |a, b| a + b);
let diff = (sum - fp("1")).abs();
assert!(diff < tight(), "softmax sum = {}, expected 1.0", sum);
}
#[test]
fn test_softmax_monotone() {
let scores = vec![fp("1"), fp("2"), fp("3")];
let result = softmax(&scores).unwrap();
assert!(result[0] < result[1], "softmax not monotone: {} >= {}", result[0], result[1]);
assert!(result[1] < result[2], "softmax not monotone: {} >= {}", result[1], result[2]);
}
#[test]
fn test_rms_norm_factor_constant() {
let c = fp("2");
let eps = fp("0.000001");
let vals = vec![c; 4];
let factor = rms_norm_factor(&vals, eps).unwrap();
let diff = (factor - fp("0.5")).abs();
assert!(diff < loose(), "rms_norm_factor = {}, expected ~0.5", factor);
}
#[test]
fn test_silu_deep_negative_is_zero() {
for s in ["-30", "-40", "-70", "-100"] {
let v = silu(fp(s));
assert!(v.abs() < tight(), "silu({}) = {}, expected ~0", s, v);
}
}
#[test]
fn test_silu_zero() {
let result = silu(FixedPoint::ZERO);
assert!(result.abs() < tight(), "silu(0) = {}, expected 0", result);
}
#[test]
fn test_silu_positive() {
let x = fp("10");
let result = silu(x);
let diff = (result - x).abs();
assert!(diff < loose(), "silu(10) = {}, expected ~10", result);
}
#[test]
fn test_silu_negative() {
let result = silu(fp("-10"));
assert!(result.abs() < loose(), "silu(-10) = {}, expected ~0", result);
}
#[test]
fn test_softmax_mix_one_hot() {
let scores = vec![fp("20"), fp("0"), fp("0")];
let rows = [
vec![fp("1"), fp("2")],
vec![fp("-5"), fp("7")],
vec![fp("3"), fp("-3")],
];
let refs: Vec<&[FixedPoint]> = rows.iter().map(|r| r.as_slice()).collect();
let (out, w) = softmax_mix(&scores, &refs).unwrap();
assert!((out[0] - fp("1")).abs() < fp("0.01"), "out[0] = {}", out[0]);
assert!((out[1] - fp("2")).abs() < fp("0.01"), "out[1] = {}", out[1]);
assert!((w[0] - fp("1")).abs() < fp("0.01"), "w[0] = {}", w[0]);
}
#[test]
fn test_softmax_mix_uniform_matches_mean() {
let scores = vec![FixedPoint::ZERO; 4];
let rows = [
vec![fp("4")],
vec![fp("8")],
vec![fp("-4")],
vec![fp("0")],
];
let refs: Vec<&[FixedPoint]> = rows.iter().map(|r| r.as_slice()).collect();
let (out, _) = softmax_mix(&scores, &refs).unwrap();
assert!((out[0] - fp("2")).abs() < fp("0.01"), "out[0] = {}", out[0]);
}
#[test]
fn test_softmax_mix_agrees_with_materialized_at_short_length() {
let scores = vec![fp("1.5"), fp("0.5"), fp("-0.25"), fp("2")];
let rows = [
vec![fp("1"), fp("-2")],
vec![fp("0.5"), fp("3")],
vec![fp("-1.5"), fp("0.25")],
vec![fp("2"), fp("1")],
];
let refs: Vec<&[FixedPoint]> = rows.iter().map(|r| r.as_slice()).collect();
let (fused_out, _) = softmax_mix(&scores, &refs).unwrap();
let w = softmax(&scores).unwrap();
for d in 0..2 {
let mut acc = FixedPoint::ZERO;
for j in 0..4 {
acc = acc + w[j] * rows[j][d];
}
let diff = (fused_out[d] - acc).abs();
assert!(
diff < fp("0.01"),
"fused vs materialized dim {}: {} vs {}",
d, fused_out[d], acc
);
}
}
#[test]
fn test_softmax_mix_survives_below_storage_floor() {
let n = 3000;
let scores = vec![FixedPoint::ZERO; n];
let rows: Vec<Vec<FixedPoint>> = (0..n)
.map(|j| vec![if j % 2 == 0 { fp("2") } else { fp("4") }])
.collect();
let refs: Vec<&[FixedPoint]> = rows.iter().map(|r| r.as_slice()).collect();
let (out, _) = softmax_mix(&scores, &refs).unwrap();
assert!(
(out[0] - fp("3")).abs() < fp("0.05"),
"fused mix should recover mean 3.0, got {}",
out[0]
);
let w = softmax(&scores).unwrap();
let mut materialized = FixedPoint::ZERO;
for j in 0..n {
materialized = materialized + w[j] * rows[j][0];
}
if materialized.abs() < fp("0.5") {
assert!(
(out[0] - fp("3")).abs() < (out[0] - materialized).abs(),
"fused mix ({}) should beat the collapsed materialized mix ({})",
out[0],
materialized
);
}
}
fn softmax_mix_reference(
scores: &[FixedPoint],
values: &[&[FixedPoint]],
) -> Result<(Vec<FixedPoint>, Vec<FixedPoint>), OverflowDetected> {
assert_eq!(
scores.len(),
values.len(),
"softmax_mix: scores/values length mismatch"
);
if scores.is_empty() {
return Ok((vec![], vec![]));
}
let dim = values[0].len();
let mut max_raw = scores[0].raw();
for s in &scores[1..] {
if s.raw() > max_raw {
max_raw = s.raw();
}
}
let max_compute = FixedPoint::upscale_to_compute(max_raw);
let mut exp_values: Vec<ComputeStorage> = Vec::with_capacity(scores.len());
let mut sum = compute_zero::<BinaryStorage, FRAC_BITS>();
for s in scores {
let shifted = FixedPoint::compute_subtract(FixedPoint::upscale_to_compute(s.raw()), max_compute);
let e = FixedPoint::exp_at_compute_tier(shifted);
sum = FixedPoint::compute_checked_add(sum, e)?;
exp_values.push(e);
}
if FixedPoint::compute_is_zero(&sum) {
return Err(OverflowDetected::DivisionByZero);
}
let mut num: Vec<ComputeStorage> = vec![compute_zero::<BinaryStorage, FRAC_BITS>(); dim];
for (j, v) in values.iter().enumerate() {
assert_eq!(
v.len(),
dim,
"softmax_mix: value row {j} has length {}, expected {dim}",
v.len()
);
let e = exp_values[j];
for d in 0..dim {
num[d] = FixedPoint::compute_checked_add(num[d], FixedPoint::compute_multiply(e, FixedPoint::upscale_to_compute(v[d].raw())))?;
}
}
let mut out = Vec::with_capacity(dim);
for n in &num {
out.push(FixedPoint::from_raw(FixedPoint::downscale_to_storage(FixedPoint::compute_divide(*n, sum)?)?));
}
let mut weights = Vec::with_capacity(scores.len());
for e in &exp_values {
weights.push(FixedPoint::from_raw(FixedPoint::downscale_to_storage(FixedPoint::compute_divide(*e, sum)?)?));
}
Ok((out, weights))
}
struct Lcg(u64);
impl Lcg {
fn next(&mut self) -> u64 {
self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
self.0 >> 33
}
fn value(&mut self, amp: i32) -> FixedPoint {
assert!(amp <= 8000);
let k = (self.next() % (2 * amp as u64 + 1)) as i32 - amp;
FixedPoint::from_int(k / 64) + FixedPoint::from_int(k % 64) / FixedPoint::from_int(64)
}
}
fn assert_mix_variants_equal_reference(scores: &[FixedPoint], rows: &[Vec<FixedPoint>], dim: usize) {
let refs: Vec<&[FixedPoint]> = rows.iter().map(|r| r.as_slice()).collect();
let flat: Vec<FixedPoint> = rows.iter().flat_map(|r| r.iter().copied()).collect();
let want = softmax_mix_reference(scores, &refs);
assert_eq!(softmax_mix(scores, &refs), want);
assert_eq!(softmax_mix_flat(scores, &flat, dim), want);
let want_out = want.map(|(out, _)| out);
assert_eq!(softmax_mix_values(scores, &refs), want_out);
assert_eq!(softmax_mix_flat_values(scores, &flat, dim), want_out);
}
#[test]
fn test_softmax_mix_variants_equal_the_checked_reference() {
let mut rng = Lcg(17);
for &dim in &[1usize, 3, 8, 64] {
for &n in &[1usize, 2, 3, 15, 16, 17, 100, 600] {
for & in &[2i32, 200, 6000] {
let scores: Vec<FixedPoint> = (0..n).map(|_| rng.value(amp)).collect();
let rows: Vec<Vec<FixedPoint>> = (0..n).map(|_| (0..dim).map(|_| rng.value(6000)).collect()).collect();
assert_mix_variants_equal_reference(&scores, &rows, dim);
}
}
}
assert_eq!(softmax_mix_values::<BinaryStorage, FRAC_BITS>(&[], &[]), Ok(vec![]));
assert_eq!(softmax_mix_flat::<BinaryStorage, FRAC_BITS>(&[], &[], 4), Ok((vec![], vec![])));
}
on_words!($W in [w32] {
#[test]
fn test_softmax_mix_i64_path_equals_reference_at_extremes() {
let mut rng = Lcg(99);
let mut raw = |amp: i64| FixedPoint::from_raw(((rng.next() as i64) % (2 * amp + 1) - amp) as i32);
for &(n, dim, score_amp, value_amp) in &[
(1usize, 1usize, 1i64 << 30, i32::MAX as i64),
(2, 1, 3, i32::MAX as i64),
(17, 5, 1 << 12, i32::MAX as i64),
(2049, 1, 1 << 6, i32::MAX as i64 / 4),
(2049, 8, 1 << 16, 1 << 20),
(4097, 2, 2, i32::MAX as i64),
(70_000, 1, 1, 1 << 30),
] {
let scores: Vec<FixedPoint> = (0..n).map(|_| raw(score_amp)).collect();
let rows: Vec<Vec<FixedPoint>> = (0..n).map(|_| (0..dim).map(|_| raw(value_amp)).collect()).collect();
assert_mix_variants_equal_reference(&scores, &rows, dim);
}
let n = 300;
let scores = vec![FixedPoint::ZERO; n];
let rows = vec![vec![FixedPoint::from_raw(i32::MIN); 3]; n];
assert_mix_variants_equal_reference(&scores, &rows, 3);
}
});
#[test]
fn test_dot_many_equals_dot_per_key() {
let mut rng = Lcg(5);
for &dim in &[1usize, 3, 15, 16, 17, 64, 128] {
for &keys in &[0usize, 1, 2, 9] {
let query: Vec<FixedPoint> = (0..dim).map(|_| rng.value(64)).collect();
let flat: Vec<FixedPoint> = (0..dim * keys).map(|_| rng.value(64)).collect();
let got = dot_many(&query, &flat, dim);
assert_eq!(got.len(), keys);
for k in 0..keys {
assert_eq!(got[k], dot(&query, &flat[k * dim..(k + 1) * dim]), "dim {dim} key {k}");
}
}
}
}
on_words!($W in [w32] {
#[test]
fn test_bounded_dot_equals_checked_dot_on_raws() {
let mut rng = Lcg(42);
let mut raws = |n: usize, amp: i64| -> Vec<FixedPoint> {
(0..n).map(|_| FixedPoint::from_raw(((rng.next() as i64) % (2 * amp + 1) - amp) as i32)).collect()
};
let exact = |a: &[FixedPoint], b: &[FixedPoint]| -> Option<i32> {
let acc: i128 = a.iter().zip(b).map(|(x, y)| x.raw() as i128 * y.raw() as i128).sum();
let r = (acc >> FRAC_BITS) + ((acc >> (FRAC_BITS - 1)) & 1);
i32::try_from(r).ok()
};
for case in 0..3000usize {
let n = [1usize, 7, 16, 17, 64, 128, 256, 1000][case % 8];
let (qa, ka) = (if case % 7 == 0 { 1i64 << 25 } else { 1 << 14 }, if case % 11 == 0 { 1i64 << 20 } else { 1 << 12 });
let (q, k) = (raws(n, qa), raws(n, ka));
let want = exact(&q, &k);
let got = std::panic::catch_unwind(|| dot(&q, &k).raw()).ok();
assert_eq!(got, want, "case {case} n {n}");
let many = std::panic::catch_unwind(|| dot_many(&q, &k, n)[0].raw()).ok();
assert_eq!(many, want, "dot_many case {case} n {n}");
}
let half = 1i32 << (FRAC_BITS - 1);
for target in [half, -half, 3 * half, -3 * half, half - 1, -half - 1] {
let mut q = vec![FixedPoint::ZERO; 32];
let mut k = vec![FixedPoint::ZERO; 32];
q[5] = FixedPoint::from_raw(target);
k[5] = FixedPoint::from_raw(1);
assert_eq!(Some(dot(&q, &k).raw()), exact(&q, &k), "tie {target}");
}
let big = vec![FixedPoint::from_raw(i32::MIN); 64];
assert!(std::panic::catch_unwind(|| dot(&big, &big)).is_err());
assert!(std::panic::catch_unwind(|| dot_many(&big, &big, 64)).is_err());
}
});
#[test]
fn test_rms_norm_one_rounding() {
let one = FixedPoint::one();
let half = one / FixedPoint::from_int(2);
let two = FixedPoint::from_int(2);
let w = [half, -(one + half), two, -one];
assert_eq!(rms_norm(&[two, two, -two, two], &w, 0), Ok(vec![half, -(one + half), -two, -one]));
let mut unit = one;
for _ in 0..FRAC_BITS { unit = unit / two; }
let tiny = [unit, FixedPoint::ZERO, FixedPoint::ZERO, FixedPoint::ZERO];
assert_eq!(rms_norm(&tiny, &[one; 4], 0), Ok(vec![two, FixedPoint::ZERO, FixedPoint::ZERO, FixedPoint::ZERO]));
let mut rng = Lcg(8);
let eps = 1i128 << 44; for &n in &[1usize, 2, 7, 64, 300] {
let x: Vec<FixedPoint> = (0..n).map(|_| rng.value(200)).collect();
let w: Vec<FixedPoint> = (0..n).map(|_| rng.value(128)).collect();
let want = rms_norm(&x, &w, eps).unwrap();
let mut y = x.clone();
assert_eq!(rms_norm_in_place(&mut y, &w, eps), Ok(()));
assert_eq!(y, want);
}
assert_eq!(rms_norm(&[FixedPoint::ZERO; 4], &[one; 4], eps), Ok(vec![FixedPoint::ZERO; 4]));
assert_eq!(rms_norm(&[FixedPoint::ZERO; 4], &[one; 4], 0), Err(OverflowDetected::DivisionByZero));
assert_eq!(rms_norm::<BinaryStorage, FRAC_BITS>(&[], &[], eps), Err(OverflowDetected::DivisionByZero));
assert_eq!(rms_norm(&[FixedPoint::ZERO; 4], &[one; 4], -eps), Err(OverflowDetected::DomainError));
let mut big = one;
while let Ok(next) = big.try_add(big) { big = next; }
assert_eq!(rms_norm(&tiny, &[big, one, one, one], 0), Err(OverflowDetected::TierOverflow));
}
}
};
}
for_each_format!(suite);
}