diff --git a/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs b/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs index 03d042192d..98982c0ce4 100644 --- a/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs +++ b/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs @@ -108,13 +108,58 @@ use onnx_runtime_ep_api::{Result, TensorMut}; /// an in-place `y = act(y)` node where the ranges do overlap — falls back to the /// owned buffer, which is exactly the situation `write_dense_f32_narrow` exists /// to handle. Correctness never depends on which arm runs. +/// +/// `f` must be a pure elementwise map: `y[i]` may depend on `x[i]` but on no +/// other index. The narrow-output arm below evaluates it in chunks, so a +/// closure that reads across lanes — a row-broadcast bias, say — would see a +/// different window per chunk. Those callers exist ([`write_mapped_reading`]'s +/// `also_read` cases), and declaring their extra range is what keeps them off +/// the chunked path; the gate is structural rather than a convention. pub(crate) fn write_mapped(op: &str, out: &mut TensorMut, x: &[f32], f: F) -> Result<()> where - F: FnOnce(&[f32], &mut [f32]), + F: Fn(&[f32], &mut [f32]), { write_mapped_reading(op, out, x, &[], f) } +/// Bulk `f32` → 16-bit narrow, as selected by the output's dtype. +/// +/// # Safety +/// The caller must have confirmed the corresponding ISA is available and pass +/// two slices of equal length. +#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] +type NarrowFn = unsafe fn(&[f32], &mut [u16]); + +/// End of the fused-narrow chunk starting at `start`, for a tensor of `n`. +/// +/// Two properties matter, and both are about staying bit-identical to a single +/// whole-tensor pass rather than about speed: +/// +/// * **No short chunk.** `dispatch!` picks the scalar path for any slice below +/// [`SIMD_MIN_LEN`], and scalar and vector are not bit-identical for the +/// approximated kernels. A tensor of `NARROW_CHUNK + 16` elements split +/// naively would run its last 16 elements scalar where one pass would have +/// run them vectorised. So a trailing remainder below `SIMD_MIN_LEN` is +/// absorbed into the chunk before it instead of becoming its own. +/// * **Lane-aligned boundaries.** `NARROW_CHUNK` is a multiple of 8, so every +/// boundary lands where the 8-wide vector run would have been anyway and the +/// scalar remainder is the same `n % 8` elements either way. +#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] +fn chunk_end(start: usize, n: usize) -> usize { + let end = (start + NARROW_CHUNK).min(n); + if n - end < SIMD_MIN_LEN { n } else { end } +} + +/// Elements of `f32` scratch held between the kernel and the narrow when the +/// two are fused. +/// +/// 8192 f32 is 32 KiB, so a chunk's intermediate is still in L1/L2 when the +/// narrow reads it back. The unfused path allocated one buffer the size of the +/// whole tensor and made three passes over it — widen, compute, narrow — which +/// on a 1 Mi tensor is 4 MB streamed twice more than it needs to be. +#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] +const NARROW_CHUNK: usize = 8192; + /// [`write_mapped`] for closures that read a slice *besides* `x`. /// /// The disjointness check has to cover every buffer the closure still reads @@ -134,7 +179,7 @@ pub(crate) fn write_mapped_reading( f: F, ) -> Result<()> where - F: FnOnce(&[f32], &mut [f32]), + F: Fn(&[f32], &mut [f32]), { let n = x.len(); let eligible = if also_read.is_empty() { @@ -158,6 +203,66 @@ where f(x, dst); return Ok(()); } + // Fuse the kernel and the narrow for a contiguous 16-bit output, the shape + // every f16/bf16 activation lands in. Chunking keeps each intermediate in + // L1 instead of writing a tensor-sized `f32` buffer out to memory and + // reading all of it back. + // + // The result stays bit-identical to the unfused arm, but that takes more + // than the kernel being elementwise: `dispatch!` picks scalar or vector + // from the length of the slice it is handed, and the two are *not* + // bit-identical for the approximated kernels. So a chunk must never be + // shorter than `SIMD_MIN_LEN` unless it is the whole tensor, and every + // boundary must be a multiple of the 8-wide lane so the vector run lines + // up with where it would have fallen in one pass. `chunk_end` maintains + // both. + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + if also_read.is_empty() && n != 0 && out.is_contiguous() && out.numel() == n { + use crate::dtype::{bf16x, f16c}; + use onnx_runtime_ir::DataType; + let narrow: Option = match out.dtype { + DataType::Float16 if f16c::available() => Some(f16c::narrow), + DataType::BFloat16 if bf16x::available() => Some(bf16x::narrow), + _ => None, + }; + if let Some(narrow) = narrow { + out.validate()?; + let out_ptr = out.data_ptr_mut::(); + let out_bytes = (out_ptr as usize)..(out_ptr as usize + n * 2); + let x_bytes = slice_byte_range(x); + debug_assert!( + x_bytes.start >= out_bytes.end || out_bytes.start >= x_bytes.end, + "the fused narrow writes `out` while still reading `x`; widening \ + to f32 always allocates, so they must never overlap" + ); + // SAFETY: a validated contiguous `Float16`/`BFloat16` output + // addresses exactly `n` 2-byte elements -- the dtype match is what + // makes the element width 16 bits -- and `half::f16`/`half::bf16` + // are `repr(transparent)` over `u16`, so writing `u16` bit patterns + // is sound. `dst` is disjoint from `x`: reaching a 16-bit output + // means `x` came from `to_dense_f32_widen`, which always allocates + // a fresh `Cow::Owned` buffer when it has to widen, so no live + // borrow of `out`'s storage can be in `x`. The `debug_assert` above + // is the tripwire for a future caller that breaks that. + let dst = unsafe { std::slice::from_raw_parts_mut(out_ptr, n) }; + // The final chunk can absorb a short tail, so size for that. + let mut scratch = vec![0.0f32; (NARROW_CHUNK + SIMD_MIN_LEN).min(n)]; + // Serial for the same reason the unfused arm is: see `serial_scope`. + serial_scope(|| { + let mut start = 0; + while start < n { + let end = chunk_end(start, n); + let mid = &mut scratch[..end - start]; + f(&x[start..end], mid); + // SAFETY: `available()` confirmed the ISA above, and the + // two slices are the same length by construction. + unsafe { narrow(mid, &mut dst[start..end]) }; + start = end; + } + }); + return Ok(()); + } + } let mut y = vec![0.0f32; n]; // Serial on purpose: see `serial_scope`. This arm always ends in a full // serial pass over `y`, so splitting the kernel only costs locality. @@ -4050,6 +4155,223 @@ mod tests { } } + /// The fused narrow arm evaluates the kernel `NARROW_CHUNK` elements at a + /// time. That is only sound if chunking is unobservable, so this pins the + /// output bit-for-bit against a single-shot reference over a length that + /// spans several chunks and ends mid-chunk — the case where an off-by-one + /// in the chunk/`dst` pairing, a stale `scratch` tail, or a kernel that + /// looked across lanes would all show up. + /// An `x` at which the scalar and vector paths of `kernel` disagree *after* + /// narrowing to f16. + /// + /// A slice shorter than [`SIMD_MIN_LEN`] takes the scalar path and a longer + /// one takes the vector path; for the approximated kernels the two are not + /// bit-identical. Most of those f32 disagreements are washed out by f16 + /// rounding, so the witness has to be checked in the same domain the test + /// asserts in, or the test proves nothing. Finding it at test time — rather + /// than hard-coding a constant a coefficient change could quietly retire — + /// is what makes the chunking test below non-vacuous by construction. + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + fn scalar_vector_divergence(kernel: fn(&[f32], &mut [f32])) -> Option { + let mut short_in = [0.0f32; 8]; + let mut short_out = [0.0f32; 8]; + let mut long_in = vec![0.0f32; SIMD_MIN_LEN * 4]; + let mut long_out = vec![0.0f32; SIMD_MIN_LEN * 4]; + let mut probe = |v: f32| { + short_in.fill(v); + long_in.fill(v); + kernel(&short_in, &mut short_out); + kernel(&long_in, &mut long_out); + half::f16::from_f32(short_out[0]).to_bits() + != half::f16::from_f32(long_out[0]).to_bits() + }; + // Sweep every f32 outward from zero across the range the activation + // approximations actually bend in; f16 has resolution here, so a + // disagreement can survive the narrow. + for sign in [1.0f32, -1.0] { + let base = (0.001f32).to_bits(); + for step in 0..200_000u32 { + let v = sign * f32::from_bits(base + step); + if v.abs() > 12.0 { + break; + } + if probe(v) { + return Some(v); + } + } + } + None + } + + /// A trailing chunk shorter than [`SIMD_MIN_LEN`] would take the scalar + /// path where a single whole-tensor pass runs the same elements vectorised, + /// and the two are not bit-identical. `chunk_end` absorbs such a tail into + /// the previous chunk; this pins that the output really is unchanged, using + /// an input value where the two paths are *known* to disagree so the test + /// cannot pass by luck. + #[test] + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + fn a_short_tail_does_not_fall_off_the_vector_path() { + let mut witnesses = 0; + for (name, kernel) in [ + ("Tanh", tanh_f32_slice as fn(&[f32], &mut [f32])), + ("Erf", erf_f32_slice as fn(&[f32], &mut [f32])), + ("Sigmoid", sigmoid_f32_slice as fn(&[f32], &mut [f32])), + ] { + let Some(v) = scalar_vector_divergence(kernel) else { + continue; + }; + witnesses += 1; + for tail in [1usize, 8, 16, 31] { + let n = NARROW_CHUNK + tail; + // Only the tail carries the divergent value; the body is + // ordinary so a failure localises to the seam. + let mut src = vec![0.25f32; n]; + src[n - tail..].fill(v); + + let mut expected = vec![0.0f32; n]; + kernel(&src, &mut expected); + + let mut got = Owned::f16(&[n], &vec![0.0f32; n]); + write_mapped(name, &mut got.view_mut(), &src, kernel).unwrap(); + for (i, (g, e)) in got.to_u16_bits().iter().zip(&expected).enumerate() { + assert_eq!( + *g, + half::f16::from_f32(*e).to_bits(), + "{name} tail {tail}: fused != unfused at {i} of {n} (x={})", + src[i] + ); + } + } + } + assert!( + witnesses > 0, + "found no scalar/vector divergence at all, so this test proves nothing" + ); + } + + #[test] + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + fn fused_narrow_is_bit_identical_across_chunk_boundaries() { + // A tail shorter than `SIMD_MIN_LEN` is the dangerous case: handed to + // the kernel on its own it would take the scalar path, while a single + // whole-tensor pass runs those same elements vectorised, and the two + // are not bit-identical. 137 is a long tail, 16 and 8 are short ones, + // and 8191 leaves a one-element remainder inside a full chunk. + for tail in [137usize, 31, 16, 8, 7, 1, 8191] { + let n = NARROW_CHUNK * 2 + tail; + // Straddle the interesting parts of each kernel's domain: the tanh + // saturation knees, erf's split boundary, and sigmoid's tail, where + // the scalar and vector approximations part company. + let src: Vec = (0..n).map(|i| (i as f32 - 8000.0) * 0.0037).collect(); + + for (name, kernel) in [ + ("Tanh", tanh_f32_slice as fn(&[f32], &mut [f32])), + ("Erf", erf_f32_slice as fn(&[f32], &mut [f32])), + ("Sigmoid", sigmoid_f32_slice as fn(&[f32], &mut [f32])), + ] { + let mut expected = vec![0.0f32; n]; + kernel(&src, &mut expected); + + let mut f16_out = Owned::f16(&[n], &vec![0.0f32; n]); + write_mapped(name, &mut f16_out.view_mut(), &src, kernel).unwrap(); + for (i, (g, e)) in f16_out.to_u16_bits().iter().zip(&expected).enumerate() { + assert_eq!( + *g, + half::f16::from_f32(*e).to_bits(), + "{name} f16 differs at {i} of {n} (tail {tail}, x={})", + src[i] + ); + } + + let mut bf16_out = Owned::bf16(&[n], &vec![0.0f32; n]); + write_mapped(name, &mut bf16_out.view_mut(), &src, kernel).unwrap(); + for (i, (g, e)) in bf16_out.to_u16_bits().iter().zip(&expected).enumerate() { + assert_eq!( + *g, + half::bf16::from_f32(*e).to_bits(), + "{name} bf16 differs at {i} of {n} (tail {tail}, x={})", + src[i] + ); + } + } + } + } + + /// `chunk_end` is what keeps the kernel off the scalar path for a short + /// trailing remainder. Pin the two properties the bit-identity argument + /// rests on for every tensor size around a chunk boundary: no chunk is + /// both short and not the whole tensor, and no boundary falls off an + /// 8-wide lane. + #[test] + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + fn chunk_end_never_emits_a_short_trailing_chunk() { + let sizes = (1..80usize) + .chain([NARROW_CHUNK - 1, NARROW_CHUNK, NARROW_CHUNK + 1]) + .chain((0..80).map(|i| NARROW_CHUNK + i)) + .chain((0..80).map(|i| NARROW_CHUNK * 3 + i)); + for n in sizes { + let mut start = 0; + let mut chunks = 0; + while start < n { + let end = chunk_end(start, n); + assert!(end > start, "n={n} made no progress at {start}"); + let len = end - start; + assert!( + len >= SIMD_MIN_LEN || (start == 0 && end == n), + "n={n} chunk [{start},{end}) is short but not the whole tensor" + ); + assert!( + end == n || end.is_multiple_of(8), + "n={n} boundary {end} is not lane-aligned" + ); + assert!( + len <= NARROW_CHUNK + SIMD_MIN_LEN, + "n={n} chunk [{start},{end}) overruns the scratch" + ); + start = end; + chunks += 1; + assert!(chunks < 16, "n={n} did not terminate"); + } + } + } + + /// A closure that reads across lanes must never reach the chunked arm. + /// `BiasGelu`/`FastGelu` declare their bias through `also_read`, and that + /// declaration is the only thing keeping them on the whole-tensor path — + /// so this asserts the row-broadcast result still matches a single-shot + /// reference when the output is f16 and the tensor spans several chunks + /// with a row width that never divides one. + #[test] + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + fn a_row_broadcast_closure_stays_off_the_chunked_arm() { + let width = 61; + let n = width * 300; + assert!(n > NARROW_CHUNK && !NARROW_CHUNK.is_multiple_of(width)); + let src: Vec = (0..n).map(|i| (i as f32 - 9000.0) * 0.0021).collect(); + let bias: Vec = (0..width).map(|i| (i as f32 - 30.0) * 0.05).collect(); + + let mut expected = vec![0.0f32; n]; + erf_gelu_bias_f32_slice(&src, &bias, width, &mut expected); + + let mut out = Owned::f16(&[n], &vec![0.0f32; n]); + write_mapped_reading( + "BiasGelu", + &mut out.view_mut(), + &src, + &[crate::dtype::slice_byte_range(&bias)], + |x, y| erf_gelu_bias_f32_slice(x, &bias, width, y), + ) + .unwrap(); + for (i, (g, e)) in out.to_u16_bits().iter().zip(&expected).enumerate() { + assert_eq!( + *g, + half::f16::from_f32(*e).to_bits(), + "row-broadcast output differs at {i}" + ); + } + } + /// Fusing FastGelu's bias into the vector kernel must be bit-identical to /// materialising `x + bias` and mapping over it — that equivalence is the /// whole justification for the fused path.