diff --git a/core/src/ops/nn/rms_norm.rs b/core/src/ops/nn/rms_norm.rs index c6ee9fe3ec..63b777e9c9 100644 --- a/core/src/ops/nn/rms_norm.rs +++ b/core/src/ops/nn/rms_norm.rs @@ -41,28 +41,32 @@ impl EvalOp for RmsNorm { && self.axis == input.rank() - 1 { let eps_f32: f32 = self.eps.cast_to_scalar::()?; - let already_f32 = in_dt == DatumType::F32; - let mut buf = if already_f32 { - input.into_tensor() - } else { - input.cast_to::()?.into_owned() - }; - let row_len = buf.shape()[self.axis]; + let row_len = input.shape()[self.axis]; + let mut buf = input.into_tensor(); if row_len > 0 { - let data = unsafe { buf.as_slice_mut_unchecked::() }; let rms_norm = &tract_linalg::ops().rms_norm_f32; - let total = data.len(); - tract_linalg::multithread::par_chunks_mut(data, row_len, total, |_, chunk| { - for row in chunk.chunks_mut(row_len) { - rms_norm(row, eps_f32); - } - Ok(()) - })?; + if in_dt == DatumType::F32 { + let data = unsafe { buf.as_slice_mut_unchecked::() }; + let total = data.len(); + tract_linalg::multithread::par_chunks_mut(data, row_len, total, |_, chunk| { + for row in chunk.chunks_mut(row_len) { + rms_norm(row, eps_f32); + } + Ok(()) + })?; + } else { + let rms_norm = &tract_linalg::ops().rms_norm_f16; + let data = unsafe { buf.as_slice_mut_unchecked::() }; + let total = data.len(); + tract_linalg::multithread::par_chunks_mut(data, row_len, total, |_, chunk| { + for row in chunk.chunks_mut(row_len) { + rms_norm(row, eps_f32); + } + Ok(()) + })?; + } } - if already_f32 { - return Ok(tvec![buf.into_tvalue()]); - } - return Ok(tvec![buf.cast_to_dt(in_dt)?.into_owned().into()]); + return Ok(tvec![buf.into_tvalue()]); } // Slow path: original 4-call composition (kept for non-contiguous axes). diff --git a/linalg/src/arm64.rs b/linalg/src/arm64.rs index 211b7be997..466399a3e3 100644 --- a/linalg/src/arm64.rs +++ b/linalg/src/arm64.rs @@ -497,6 +497,7 @@ pub fn plug(ops: &mut Ops) { ops.sum_f32 = Box::new(|| arm64simd_sum_f32_16n::red()); ops.mul_by_scalar_f32 = Box::new(|| arm64simd_mul_by_scalar_f32_16n::ew()); ops.softmax2_fastcompact_f32 = Box::new(|| arm64simd_softmax2_fastcompact_f32_16n::red()); + ops.rms_norm_f16 = Box::new(arm64simd_rms_norm_f16); ops.rms_norm_f32 = Box::new(arm64simd_rms_norm_f32); #[cfg(not(feature = "no_fp16"))] if has_fp16() { diff --git a/linalg/src/arm64/arm64simd.rs b/linalg/src/arm64/arm64simd.rs index 0f6fc17e4e..cd47dfeb6c 100644 --- a/linalg/src/arm64/arm64simd.rs +++ b/linalg/src/arm64/arm64simd.rs @@ -23,6 +23,7 @@ pub use hardswish::arm64simd_hardswish_f32_8n; pub use leaky_relu::arm64simd_leaky_relu_f32_8n; pub use max::arm64simd_max_f32_16n; pub use min::arm64simd_min_f32_16n; +pub use rms_norm::rms_norm_f16 as arm64simd_rms_norm_f16; pub use rms_norm::rms_norm_f32 as arm64simd_rms_norm_f32; pub use silu::arm64simd_silu_f32_4n; pub use silu_fused::arm64simd_silu_f32_4n_fused; diff --git a/linalg/src/arm64/arm64simd/rms_norm.rs b/linalg/src/arm64/arm64simd/rms_norm.rs index 5f25299468..762db214b2 100644 --- a/linalg/src/arm64/arm64simd/rms_norm.rs +++ b/linalg/src/arm64/arm64simd/rms_norm.rs @@ -13,6 +13,8 @@ // dispatcher in `core::ops::nn::RmsNorm::eval` is already arch-neutral and // will pick this up automatically. +use tract_data::internal::f16; + #[target_feature(enable = "neon")] unsafe fn rms_norm_f32_inner(buf: &mut [f32], eps: f32) { use std::arch::aarch64::*; @@ -102,6 +104,108 @@ pub fn rms_norm_f32(buf: &mut [f32], eps: f32) { unsafe { rms_norm_f32_inner(buf, eps) } } +// f16 rows, arithmetic unchanged from `rms_norm_f32_inner`: FCVTL widens each +// 16-element group into the same four f32 registers before the same FMLA chain, +// and FCVTN rounds to nearest-even on the way out, matching `f16::from_f32`. +#[target_feature(enable = "neon")] +unsafe fn rms_norm_f16_inner(buf: &mut [f16], eps: f32) { + use std::arch::aarch64::*; + let n = buf.len(); + let chunks = n / 16; + let tail_start = chunks * 16; + let ptr = buf.as_mut_ptr(); + + let mut sum_sq: f32 = 0.0; + if chunks > 0 { + let p = ptr; + let c = chunks; + let mut sum_v: float32x4_t = vdupq_n_f32(0.0); + unsafe { + std::arch::asm!(" + movi v1.4s, 0 + movi v2.4s, 0 + movi v3.4s, 0 + 2: + ld1 {{v16.8h, v17.8h}}, [{p}], 32 + fcvtl v4.4s, v16.4h + fcvtl2 v5.4s, v16.8h + fcvtl v6.4s, v17.4h + fcvtl2 v7.4s, v17.8h + fmla v0.4s, v4.4s, v4.4s + fmla v1.4s, v5.4s, v5.4s + fmla v2.4s, v6.4s, v6.4s + fmla v3.4s, v7.4s, v7.4s + subs {c}, {c}, 1 + bne 2b + fadd v0.4s, v0.4s, v1.4s + fadd v2.4s, v2.4s, v3.4s + fadd v0.4s, v0.4s, v2.4s + ", + p = inout(reg) p => _, + c = inout(reg) c => _, + inout("v0") sum_v, + out("v1") _, out("v2") _, out("v3") _, + out("v4") _, out("v5") _, out("v6") _, out("v7") _, + out("v16") _, out("v17") _, + ); + } + sum_sq = vaddvq_f32(sum_v); + } + for i in tail_start..n { + let x = unsafe { buf.get_unchecked(i).to_f32() }; + sum_sq += x * x; + } + + let mean_sq = sum_sq / (n as f32); + let inv_std = (mean_sq + eps).sqrt().recip(); + + if chunks > 0 { + let p = ptr; + let c = chunks; + let inv_v: float32x4_t = vdupq_n_f32(inv_std); + unsafe { + std::arch::asm!(" + 2: + ld1 {{v16.8h, v17.8h}}, [{p}] + fcvtl v4.4s, v16.4h + fcvtl2 v5.4s, v16.8h + fcvtl v6.4s, v17.4h + fcvtl2 v7.4s, v17.8h + fmul v4.4s, v4.4s, v0.4s + fmul v5.4s, v5.4s, v0.4s + fmul v6.4s, v6.4s, v0.4s + fmul v7.4s, v7.4s, v0.4s + fcvtn v16.4h, v4.4s + fcvtn2 v16.8h, v5.4s + fcvtn v17.4h, v6.4s + fcvtn2 v17.8h, v7.4s + st1 {{v16.8h, v17.8h}}, [{p}], 32 + subs {c}, {c}, 1 + bne 2b + ", + p = inout(reg) p => _, + c = inout(reg) c => _, + in("v0") inv_v, + out("v4") _, out("v5") _, out("v6") _, out("v7") _, + out("v16") _, out("v17") _, + ); + } + } + for i in tail_start..n { + unsafe { + let x = buf.get_unchecked(i).to_f32(); + *buf.get_unchecked_mut(i) = f16::from_f32(x * inv_std); + } + } +} + +pub fn rms_norm_f16(buf: &mut [f16], eps: f32) { + if buf.is_empty() { + return; + } + unsafe { rms_norm_f16_inner(buf, eps) } +} + #[cfg(test)] mod tests { use super::*; @@ -160,5 +264,26 @@ mod tests { let mut x: Vec = vec![]; rms_norm_f32(&mut x, 1e-5); assert!(x.is_empty()); + let mut h: Vec = vec![]; + rms_norm_f16(&mut h, 1e-5); + assert!(h.is_empty()); + } + + #[test] + fn f16_matches_widening_the_row_to_f32() { + for len in [1usize, 7, 15, 16, 17, 33, 64, 129, 1024, 4096] { + let row: Vec = + (0..len).map(|i| f16::from_f32((i as f32 * 0.13).sin() * 5.0)).collect(); + + let mut got = row.clone(); + rms_norm_f16(&mut got, 1e-5); + + let mut widened: Vec = row.iter().map(|x| x.to_f32()).collect(); + rms_norm_f32(&mut widened, 1e-5); + let want: Vec = widened.into_iter().map(f16::from_f32).collect(); + + let mismatch = got.iter().zip(&want).position(|(a, b)| a.to_bits() != b.to_bits()); + assert_eq!(mismatch, None, "len {len}"); + } } } diff --git a/linalg/src/generic/rms_norm.rs b/linalg/src/generic/rms_norm.rs index 40eed4dbfb..2ccc281410 100644 --- a/linalg/src/generic/rms_norm.rs +++ b/linalg/src/generic/rms_norm.rs @@ -1,3 +1,5 @@ +use tract_data::internal::f16; + /// Generic scalar reference implementation of fused row-wise RmsNorm. /// out_i = x_i * rsqrt(mean(x_i²) + eps) /// @@ -17,6 +19,28 @@ pub fn rms_norm_f32(buf: &mut [f32], eps: f32) { } } +/// f16 counterpart of [`rms_norm_f32`], for hosts with no native f16 kernel. +/// +/// Widens one row at a time into a reused thread-local buffer and defers to +/// whichever f32 kernel this host registered, so an arch that has a fast f32 +/// kernel but no f16 one keeps its arithmetic and its speed. f16 -> f32 is +/// exact, so the result matches widening the row in the caller. +pub fn rms_norm_f16(buf: &mut [f16], eps: f32) { + if buf.is_empty() { + return; + } + thread_local! { + static WIDE: std::cell::RefCell> = const { std::cell::RefCell::new(Vec::new()) }; + } + WIDE.with(|wide| { + let mut wide = wide.borrow_mut(); + wide.clear(); + wide.extend(buf.iter().map(|x| x.to_f32())); + (crate::ops().rms_norm_f32)(&mut wide, eps); + buf.iter_mut().zip(wide.iter()).for_each(|(h, f)| *h = f16::from_f32(*f)); + }) +} + #[cfg(test)] mod tests { use super::*; diff --git a/linalg/src/lib.rs b/linalg/src/lib.rs index feba41da18..c8a86d2992 100644 --- a/linalg/src/lib.rs +++ b/linalg/src/lib.rs @@ -115,6 +115,7 @@ pub struct Ops { /// Replaces a 4-call composition (MeanOfSquares + Add + Rsqrt + Mul) with /// a single 2-pass kernel. Called once per row by `core::ops::nn::RmsNorm` /// when the input is f32 and the axis is the last (contiguous) one. + pub rms_norm_f16: Box, pub rms_norm_f32: Box, } @@ -253,6 +254,7 @@ pub fn generic() -> Ops { */ softmax2_fastcompact_f16: Box::new(|| generic::reduce::softmax_l2::HSoftMaxL2::red()), softmax2_fastcompact_f32: Box::new(|| generic::reduce::softmax_l2::SSoftMaxL2::red()), + rms_norm_f16: Box::new(generic::rms_norm::rms_norm_f16), rms_norm_f32: Box::new(generic::rms_norm::rms_norm_f32), }; crate::generic::mmm::plug(&mut ops);