Repository navigation
perf(cpu-ep): give x86 Relu and Clip the vector kernel they never had - #1355
Closed
justinchuby wants to merge 6 commits into
Closed
justinchuby wants to merge 6 commits into
justinchuby wants to merge 6 commits into
Conversation
`relu_f32_simd` and `clip_f32_simd` dispatch to a hand-written NEON
kernel on aarch64 and to a scalar loop everywhere else. Every other
activation in this crate got an AVX2 kernel; these two were left behind,
and it shows: `Relu` at 1 Mi is the only large-tensor case in the
benchmark grid still slower than ORT's own CPU EP, which runs MLAS's
vectorised `MlasComputeRelu` there.
Both are now eight lanes wide on x86-64, four vectors per iteration,
with the usual 8-wide remainder and scalar tail. No new dependency, no
feature gate, no MLAS: `is_x86_feature_detected!("avx2")` at the top of
the dispatcher, same shape as `simd_activations`.
The operand order is the whole correctness argument. `vmaxps` is not a
general `max` — Intel defines it as `IF SRC1 > SRC2 THEN SRC1 ELSE SRC2`
and it returns SRC2 on an unordered compare — so with the zero vector as
SRC1 and the data as SRC2 it reproduces `if x < 0.0 { 0.0 } else { x }`
exactly, NaN payloads and `-0.0` included. Written the other way round,
the way anyone would write it first, both become `+0.0`. Clip is the
same argument twice, for `vmaxps` against the lower bound and `vminps`
against the upper.
Three tests compare against the scalar form **by bit pattern**, over
both zeros, both infinities, quiet/signalling NaNs with distinct
payloads, the subnormal boundary, and every length from 0 to 114 so the
32-wide body, the 8-wide remainder and the tail each see special values;
Clip additionally over eight bound pairs including inverted and NaN
bounds. One test drives the dispatcher rather than the kernel, so a
wrong `#[cfg]` arm cannot pass.
Falsifier, run: swapping the `_mm256_max_ps` operands fails both Relu
tests at lane 1 with `0.0 (0x00000000)` against `-0.0 (0x80000000)`.
1420 `-p onnx-runtime-ep-cpu --lib` tests pass, clippy and fmt clean.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
`Relu` had only a 1 Mi bench case. At that size the op moves 8 MiB and runs at ~63 GB/s on one core — it is bandwidth-bound, so it reports almost nothing about the arithmetic. 4 Ki fits in L1, is the size a decode step actually sees, and is where every other activation in the grid is already measured. The first thing it says is worth knowing: on `origin/main` Relu at 4 Ki is 1.64x ORT, the **worst** ratio in the 4 Ki grid, and the AVX2 kernel in the previous commit does not move it (1.642 -> 1.638, 2/5 rounds). Relu is the cheapest op in the set, so ORT's own time is smallest and our fixed per-`Run` cost is the largest share of it. That is a dispatch problem, not a kernel one, and it is the number to watch as #1244/#1246 and their successors land. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
From the independent review of #1355. Neither is a correctness defect; both are gaps the review was right to name. - The AVX2 arms were gated `target_arch = "x86_64"` while the sibling f16c/bf16 staging in the same file uses `any(x86, x86_64)`. 32-bit x86 can have AVX2, so relu and clip were silently falling to the scalar loop there for no reason. Widened, with the matching `std::arch::x86` import and test gate. - `relu_f32_avx2_is_bit_identical_to_the_scalar_form` swept `0..src.len()`, which stops one short of the full 114-element slice. `0..=` now, symmetric with the clip test. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
A `git add -A` while applying the review nits swept up the A/B driver, its two result CSVs, the summariser and a scratch copy of the PR body. None of them belong in the tree. Removed from the index; they stay on disk as untracked working files. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This was referenced Aug 19, 2026
🔴 Benchmark Regression DetectedComparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).
Visual flags: Host infoWhat this cannot catch
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1355 +/- ##
===========================================
- Coverage 82.10% 80.80% -1.31%
===========================================
Files 12 377 +365
Lines 5471 167211 +161740
Branches 5471 167211 +161740
===========================================
+ Hits 4492 135111 +130619
- Misses 780 27253 +26473
- Partials 199 4847 +4648
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
justinchuby
added a commit
that referenced
this pull request
Aug 19, 2026
#1376) ## What The f16/bf16 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 is a 4 MB buffer streamed out to memory and read back for no reason — the intermediate is dead the moment the narrow has seen it. This fuses the last two passes. When the output is a contiguous f16/bf16 tensor and the closure declares no extra reads, `write_mapped_reading` evaluates the kernel ~8192 elements at a time into a 32 KiB scratch and narrows each chunk straight into the output, so the intermediate never leaves L1. The code already carried the diagnosis and named this as the fix: > The fix for f16/bf16 is to fuse widen, compute and narrow into a single pass > per chunk so each thread keeps its slice in L2 — a different change, and one > that has to be measured on its own. This is that change, measured on its own. ## Why the chunking is bit-identical Two separate things have to hold, and the second one is subtle enough that the first version of this PR got it wrong (see the review section). **1. The closure must be a pure elementwise map.** A closure that reads across lanes would see a different window per chunk. The gate is structural rather than a convention: - the row-broadcast callers (`BiasGelu`, `FastGelu` with a bias) **must** declare their bias range through `also_read`, because that is what lets `output_direct_write_eligible` reject an aliasing output; - a non-empty `also_read` is exactly what keeps them on the whole-tensor path. A caller cannot reach the chunked arm without having already asserted it reads nothing but `x`. All five `write_mapped` (empty `also_read`) call sites were checked and are plain elementwise maps: `Gelu` ×2, `QuickGelu`, `FastGelu` without bias, and the `Tanh`/`Sqrt`/`Erf` arm of unary elementwise. **2. No chunk may be shorter than `SIMD_MIN_LEN`.** `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**. A plain `chunks(NARROW_CHUNK)` hands the last chunk whatever is left over, so a tensor of `NARROW_CHUNK + 16` elements would run its final 16 scalar where one pass runs them vectorised. `chunk_end` absorbs a trailing remainder below `SIMD_MIN_LEN` into the chunk before it, and keeps every other boundary on an 8-wide lane so the vector run lands where one pass would have put it. This matters because the module treats bit reproducibility as a **correctness** property rather than a tolerance — see the `host_chunk_len` docs and the thread-invariance tests — and bf16 results cross to the CUDA EP. ### `unsafe` The `&mut [u16]` is built under `validate()`, `is_contiguous()`, `numel() == n`, and a dtype match on `Float16`/`BFloat16` (which is what makes the element width 16 bits). Disjointness from `x` does **not** rest on access ordering — that would not make overlapping `&[f32]`/`&mut [u16]` sound anyway. It rests on `to_dense_f32_widen` always allocating a fresh buffer when it widens, so a 16-bit output cannot be aliased by its own widened input. A `debug_assert` is the tripwire for a future caller that breaks that. ## Measurement Production A/B, `plugin_path_ab_vs_plain_ort`, 1 Mi elements, **1 thread**, 60 iterations, box gated to ≥93% idle. Ratio is **ours/ORT, lower is better**. | case | ratio_p50 before | after | ours p50 before → after | |---|---|---|---| | `bench_tanh_f16_1m` | 1.408 | **1.213** | 0.752 → 0.653 ms (**−13.2%**) | | `bench_exp_f16_1m` | 1.206 | 1.206 | 0.718 → 0.718 ms (flat) | Repeatability of the headline case — seven independent rounds, the last three taken **after** the correctness fix, which did not move it: | round | ratio_p50 | ours p50 | ORT p50 (control) | |---|---|---|---| | 1 | 1.217 | 0.6458 | 0.5306 | | 2 | 1.214 | 0.6542 | 0.5389 | | 3 | 1.214 | 0.6525 | 0.5376 | | 4 | 1.212 | 0.6536 | 0.5391 | | 5 (post-fix) | 1.212 | 0.6562 | 0.5415 | | 6 (post-fix) | 1.214 | 0.6524 | 0.5372 | | 7 (post-fix) | 1.214 | 0.6590 | 0.5427 | The ORT arm moves 1.0% across the whole set against a 13.2% effect, so the result is an order of magnitude clear of its own control. `Exp` does not move because it never reaches `write_mapped` — only `Tanh`, `Sqrt` and `Erf` route through the f32 path today. Widening that routing is the obvious follow-up and is deliberately **not** in this PR. Small tensors are unaffected: the scratch is `(NARROW_CHUNK + SIMD_MIN_LEN).min(n)`, so for any `n` below a chunk it allocates exactly the same `n` floats the old path did. ### Context: why f16 was the thing to fix A full 1 Mi sweep of the activation family at 1 thread showed the native kernels are already ahead of ORT on **17 of 18 f32 cases** — `mish` 0.101, `celu` 0.135, `elu` 0.137, `selu` 0.145, `log` 0.203, `sqrt` 0.309, `erf` 0.428, `tanh` 0.444, `sigmoid` 0.451 — with `relu` (1.026) the only exception, and that one is already addressed by #1355. The two 16-bit cases were the real remaining kernel-side losses, which is what pointed here. ## Tests - `a_short_tail_does_not_fall_off_the_vector_path` — the one that guards the seam. It **searches at run time** 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 proves nothing), asserts a witness was found, and then pins fused == unfused for tails of 1, 8, 16 and 31. - `chunk_end_never_emits_a_short_trailing_chunk` — pins the decomposition itself across ~250 tensor sizes around chunk boundaries: every chunk is either ≥ `SIMD_MIN_LEN` or the whole tensor, every interior boundary is lane-aligned, and no chunk overruns the scratch. - `fused_narrow_is_bit_identical_across_chunk_boundaries` — f16 **and** bf16 against a single-shot reference across `Tanh`/`Erf`/`Sigmoid` for seven tail lengths. - `a_row_broadcast_closure_stays_off_the_chunked_arm` — a `BiasGelu` with row width 61 (never divides a chunk) over a multi-chunk tensor. Every one was verified non-vacuous by mutating production code and observing the failure. Reverting `chunk_end` to the naive form fails the seam test at index **8192** with f16 bits 5146 vs 5145. Suites, on this branch: - `onnx-runtime-ep-cpu --lib`, **debug** profile: **1451 passed, 0 failed** - `plugin_ort_e2e` with `NXRT_REQUIRE_ORT_TESTS=1`: **55 passed, 0 failed**, including `every_assigned_node_is_also_executed_by_this_ep` — assigned == executed conformance holds. - `cargo clippy -p onnx-runtime-ep-cpu --lib --all-targets`: clean. - `cargo fmt -p onnx-runtime-ep-cpu -- --check`: clean. ## Review Independently reviewed by Claude Opus 4.8 with a brief to attack the safety argument and verify every test by mutation. Verdict: **REQUEST CHANGES**, one blocker, which is the `SIMD_MIN_LEN` tail divergence described above. The reviewer reproduced it with a `Sigmoid` f16 tensor of 8208 elements (bits 3455 vs 3454) and correctly observed that the original 137-element tail could not catch it. Fixed in `887c22e51`, along with the reviewer's point that the aliasing justification was argued from access ordering rather than from the allocation that actually guarantees it. The reviewer also confirmed: all eight call sites audited and the structural `also_read` gate holds; the `FnOnce`→`Fn` relaxation is safe (every closure captures only `Copy` state, no `Cell`/RNG/lazy-init); the guard set matches `write_dense_f32_narrow`; `n == 0`, `numel() != n` and non-x86 all fall through unchanged; and small-`n` allocation is not a regression. ## Scope No ORT CPU fallback, no runtime or default MLAS — pure native x86, one function plus its tests. --- ## Re-validated on latest `main` (`e0aedd0fa`), 2026-08-19 Latest `main` merged in normally (no rebase). Re-measured paired against `main` from the same box, 1 thread pinned, 60 iters / 10 warmup, `ours_p50 / ort_p50`: | arm | reps | ours p50 (ms) | ratio | |---|---|---|---| | `main` | 3 | 0.7367 / 0.7375 / 0.7319 | 1.361 / 1.367 / 1.371 | | this PR | 5 | 0.6536 / 0.6612 / 0.6577 / 0.6532 / 0.6506 | 1.213 / 1.226 / 1.226 / 1.218 / 1.217 | **Ratio 1.367 → 1.218 (median), our own time −11%.** ORT drifted under 1% between the arms (0.534–0.541 on `main`, 0.535–0.539 here), which is the control. The distributions do not overlap. ## Cross-platform fix (`f9a97e853`) The fused arm is behind `#[cfg(x86)]`, but `NarrowFn`, `chunk_end`, `NARROW_CHUNK`, the test helper `scalar_vector_divergence` and the four tests that exercise chunk boundaries were **not**. An aarch64 build warned on every one of them, and CI builds with `-D warnings` — so this branch would have failed cross-platform CI while a local x86 `clippy` run stayed green. Now gated alongside the code they belong to: **0 warnings on both x86-64 and aarch64**. All four tests assert properties of the chunked arm, which does not exist off x86, so gating them loses no coverage. ## Validation on the merged base - `cargo fmt` clean; `cargo clippy --all-targets -D warnings` clean - **1556 `onnx-runtime-ep-cpu` tests**, debug profile (`debug_assert`s live) - **55 plugin conformance tests** (`NXRT_REQUIRE_ORT_TESTS=1`, release) - `every_assigned_node_is_also_executed_by_this_ep` and `every_fixture_loads_with_cpu_fallback_disabled` green — nothing defers, nothing falls back to the ORT CPU EP - **aarch64-unknown-linux-gnu** `--all-targets` check clean, 0 warnings --------- Co-authored-by: Resch <resch@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closed unmerged — controlled null result
This PR is correct and it does what the title says: it gives x86 a hand-written
AVX2
Relu/Clipthat is bit-identical to the scalar path on NaN, ±0 and inf.It just does not make anything faster, so it is not worth the
unsafeand themaintenance.
The measurement
Paired against latest
main(68fc9480f), 1 thread pinned, 400 iters / 50warmup, 3 reps each,
ours_p50 / ort_p50:mainbench_relu_f32_4kbench_relu_f32_1mThe distributions overlap completely. There is no effect at either size.
Why — verified in the emitted assembly, not guessed
I dumped the release assembly for
onnx-runtime-ep-cpuonmainand looked at<ReluOp as ElementwiseOp>::apply_f32_bulkdirectly. The "scalar" loopis already auto-vectorised by LLVM into
maxps xmm…— SSE, 4 floats perinstruction. This PR widens that to
vmaxps ymm…, 8 floats per instruction.Doubling the vector width buys nothing here because
Reluhas no arithmeticintensity to exploit — one compare per 8 bytes of traffic (4 in, 4 out):
irrelevant;
the fixed per-
Runoverhead tracked in Reduce cpu ep kernel dispatch overhead #1077. The kernel itself is ~0.3 µs,and 4-wide was already enough to get there.
There is no size in between where this would show up, because
Reluis memory-bound at every size large enough for dispatch not to dominate.
Disposition
Closing unmerged, same standard as #1422 and #1358: a change that adds
unsafeand a second code path has to pay for itself in a controlled measurement, and
this one does not. Recorded as a falsified direction in the #1077 ledger.
The one piece here with standalone value is the
bench_relu_f32_4kcase (theexisting grid only had the bandwidth-bound 1 Mi size, which cannot say anything
about the kernel). That can be re-landed on its own if wanted; it does not need
the AVX2 code.
Refs #1077