perf(cpu-ep): fuse the f16/bf16 activation narrow into the kernel pass

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<F>(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: 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<NarrowFn> = 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::<u16>();
+            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<f32> {
+        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<f32> = (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<f32> = (0..n).map(|i| (i as f32 - 9000.0) * 0.0021).collect();
+        let bias: Vec<f32> = (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.