From 712655295b5b8f6b9d0958f199ce8cccf3ee31a4 Mon Sep 17 00:00:00 2001 From: Resch Date: Wed, 19 Aug 2026 02:55:13 +0000 Subject: [PATCH 1/3] perf(cpu-ep): fuse the f16/bf16 activation narrow into the kernel pass The 16-bit activation path was a three-pass sandwich: widen the input into an f32 scratch the size of the whole tensor, run the kernel over all of it, then narrow all of it back. On a 1 Mi tensor that streams a 4 MB buffer to memory and reads it back for no reason -- the intermediate is dead the moment the narrow has seen it. Fuse the last two passes. When the output is a contiguous f16/bf16 tensor and the closure declares no extra reads, evaluate the kernel 8192 elements at a time into a 32 KiB scratch and narrow each chunk straight into the output, so the intermediate never leaves L1. Chunking is only sound for a pure elementwise map, and the gate for that is structural rather than a convention: the row-broadcast callers (BiasGelu, FastGelu with a bias) have to declare their bias through `also_read` so the direct-write arm can reject an aliasing output, and a non-empty `also_read` is exactly what keeps them on the whole-tensor path. Measured, 1 Mi, 1 thread, 60 iters, ours/ORT p50 (lower is better): Tanh f16 1.408 -> 1.213 (0.752 -> 0.653 ms, -13.2%) Four rounds landed on 1.217/1.214/1.214/1.212 with the ORT arm moving 1.0%, so the effect is an order of magnitude clear of the control. Exp f16 does not move because `Exp` never reaches `write_mapped` -- only Tanh, Sqrt and Erf route through the f32 path today -- which is the next thing to pull. Values are bit-identical to the unfused arm, and the new tests pin that across a length spanning several chunks and ending on a partial one. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../src/kernels/simd_activations.rs | 148 +++++++++++++++++- 1 file changed, 146 insertions(+), 2 deletions(-) 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 6429c16e88..fdd95e19fd 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,36 @@ 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. +type NarrowFn = unsafe fn(&[f32], &mut [u16]); + +/// 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. +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 +157,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 +181,43 @@ 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 values are bit-identical to the unfused arm + // because the kernel is elementwise and the narrow sees the same f32 bits. + #[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()?; + // SAFETY: a validated contiguous 16-bit float output addresses + // exactly `n` 2-byte elements, and `half::f16`/`half::bf16` are + // `repr(transparent)` over `u16`, so writing `u16` bit patterns is + // sound. Each chunk of `x` is consumed into `scratch` before the + // matching half of `dst` is written, so even an output that aliases + // `x` reads its input before it is overwritten. + let dst = unsafe { std::slice::from_raw_parts_mut(out.data_ptr_mut::(), n) }; + let mut scratch = vec![0.0f32; NARROW_CHUNK.min(n)]; + // Serial for the same reason the unfused arm is: see `serial_scope`. + serial_scope(|| { + for (src, dst) in x.chunks(NARROW_CHUNK).zip(dst.chunks_mut(NARROW_CHUNK)) { + let mid = &mut scratch[..src.len()]; + f(src, mid); + // SAFETY: `available()` confirmed the ISA above, and the + // two slices are the same length by construction. + unsafe { narrow(mid, dst) }; + } + }); + 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. @@ -4003,6 +4063,90 @@ 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. + #[test] + fn fused_narrow_is_bit_identical_across_chunk_boundaries() { + let n = NARROW_CHUNK * 2 + 137; + assert!( + !n.is_multiple_of(NARROW_CHUNK), + "the test must end on a partial chunk" + ); + // Straddle the interesting parts of each kernel's domain: the tanh + // saturation knees and erf's split boundary. + 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} (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} (x={})", + src[i] + ); + } + } + } + + /// 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] + 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. From 887c22e512e1b469d56204150477319e62a85867 Mon Sep 17 00:00:00 2001 From: Resch Date: Wed, 19 Aug 2026 03:36:04 +0000 Subject: [PATCH 2/3] fix(cpu-ep): keep a short fused-narrow tail off the scalar path Review caught a real divergence in the fused arm, and the first version of the test could not see it. `dispatch!` picks scalar or vector from the length of the slice it is handed, and for the approximated kernels the two are not bit-identical. Chunking with a plain `chunks(NARROW_CHUNK)` gives the last chunk whatever is left over, so a tensor of `NARROW_CHUNK + 16` elements ran its final 16 through the scalar path where a single whole-tensor pass would have run them vectorised. That is a silent change in emitted f16/bf16 bits versus main for every size with `numel % 8192` in `[8, 32)` -- and this module treats bit reproducibility as a correctness property, not a tolerance, because bf16 crosses to the CUDA EP. `chunk_end` now absorbs a trailing remainder shorter than `SIMD_MIN_LEN` into the chunk before it, and keeps every other boundary on an 8-wide lane so the vector run lands exactly where one pass would have put it. The original test used a 137-element tail, which is above `SIMD_MIN_LEN` and so structurally could not catch this. `a_short_tail_does_not_fall_off_the _vector_path` replaces that confidence with a witness it finds at run time: it searches for an input where the scalar and vector paths disagree *after* narrowing to f16 -- most f32 disagreements are washed out by the narrow, so checking in the f32 domain would have proved nothing -- and asserts the test found one. Reverting `chunk_end` to the naive form fails it at index 8192 with f16 bits 5146 vs 5145. Also state the real disjointness argument for the `unsafe` slice: a widened f16/bf16 input is always a fresh allocation, so it cannot alias the narrower output. Access ordering would not have made overlapping `&[f32]`/`&mut [u16]` sound. A `debug_assert` is the tripwire for a future caller that breaks it. Tanh f16 1 Mi is unchanged by the fix: 1.212 / 1.214 / 1.214. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../src/kernels/simd_activations.rs | 256 +++++++++++++++--- 1 file changed, 213 insertions(+), 43 deletions(-) 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 fdd95e19fd..93092a33cb 100644 --- a/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs +++ b/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs @@ -129,6 +129,25 @@ where /// two slices of equal length. 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. +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. /// @@ -184,8 +203,16 @@ where // 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 values are bit-identical to the unfused arm - // because the kernel is elementwise and the narrow sees the same f32 bits. + // 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}; @@ -197,22 +224,37 @@ where }; if let Some(narrow) = narrow { out.validate()?; - // SAFETY: a validated contiguous 16-bit float output addresses - // exactly `n` 2-byte elements, and `half::f16`/`half::bf16` are - // `repr(transparent)` over `u16`, so writing `u16` bit patterns is - // sound. Each chunk of `x` is consumed into `scratch` before the - // matching half of `dst` is written, so even an output that aliases - // `x` reads its input before it is overwritten. - let dst = unsafe { std::slice::from_raw_parts_mut(out.data_ptr_mut::(), n) }; - let mut scratch = vec![0.0f32; NARROW_CHUNK.min(n)]; + 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(|| { - for (src, dst) in x.chunks(NARROW_CHUNK).zip(dst.chunks_mut(NARROW_CHUNK)) { - let mid = &mut scratch[..src.len()]; - f(src, mid); + 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, dst) }; + unsafe { narrow(mid, &mut dst[start..end]) }; + start = end; } }); return Ok(()); @@ -4069,45 +4111,173 @@ mod tests { /// 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. - #[test] - fn fused_narrow_is_bit_identical_across_chunk_boundaries() { - let n = NARROW_CHUNK * 2 + 137; - assert!( - !n.is_multiple_of(NARROW_CHUNK), - "the test must end on a partial chunk" - ); - // Straddle the interesting parts of each kernel's domain: the tanh - // saturation knees and erf's split boundary. - let src: Vec = (0..n).map(|i| (i as f32 - 8000.0) * 0.0037).collect(); + /// 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. + 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] + 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 mut expected = vec![0.0f32; n]; - kernel(&src, &mut expected); + 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" + ); + } - 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} (x={})", - src[i] - ); + #[test] + 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] + ); + } } + } + } - 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} (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] + 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"); } } } From f9a97e853b5befd6be1a20675286883e1b3ff972 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 19 Aug 2026 23:40:03 +0000 Subject: [PATCH 3/3] fix(cpu-ep): gate the fused-narrow helpers and their tests to x86 The fused widen/compute/narrow arm is behind #[cfg(x86)], but NarrowFn, chunk_end, NARROW_CHUNK and the four tests that exercise chunk boundaries were not. An aarch64 build warned on all of them, and CI builds with -D warnings, so this would have failed there while a local x86 clippy run stayed green. All four tests assert properties of the chunked arm, which does not exist off x86, so gating them loses no coverage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../onnx-runtime-ep-cpu/src/kernels/simd_activations.rs | 8 ++++++++ 1 file changed, 8 insertions(+) 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 6226b1d383..98982c0ce4 100644 --- a/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs +++ b/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs @@ -127,6 +127,7 @@ where /// # 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`. @@ -143,6 +144,7 @@ type NarrowFn = unsafe fn(&[f32], &mut [u16]); /// * **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 } @@ -155,6 +157,7 @@ fn chunk_end(start: usize, n: usize) -> usize { /// 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`. @@ -4168,6 +4171,7 @@ mod tests { /// 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]; @@ -4206,6 +4210,7 @@ mod tests { /// 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 [ @@ -4246,6 +4251,7 @@ mod tests { } #[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 @@ -4298,6 +4304,7 @@ mod tests { /// 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]) @@ -4336,6 +4343,7 @@ mod tests { /// 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;