Repository navigation
feat(ep-cpu): implement ONNX STFT - #2083
Merged
Merged
Conversation
Add the opset-17 CPU kernel, reusable FFT scratch planning, strict schema validation, and shared native/plugin shape inference coverage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #2083 +/- ##
==========================================
+ Coverage 80.27% 80.93% +0.66%
==========================================
Files 422 428 +6
Lines 198030 208911 +10881
Branches 198030 208911 +10881
==========================================
+ Hits 158964 169083 +10119
- Misses 33474 34132 +658
- Partials 5592 5696 +104
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
Owner
Author
|
This merged with The failure. Why only macOS. The numerics assertion above it passes. macOS is failing this test for taking the better path. Why the assertion needs reshaping rather than a The durable version is a counter on the naive branch ( One hazard to design out while you are in there, independent of the platform bug. I have not pushed a fix: I cannot run macOS here, and a blind change to the vDSP path is worth less than you making it with the lane in front of you. Happy to take it if you would rather not — say the word. For what it is worth, the counters themselves are the right instinct; a numerics test that cannot tell which path produced the numbers is the usual way a fast path silently stops being taken. This is that instrument being one platform behind the code it watches. |
This was referenced Aug 25, 2026
Owner
Author
|
Took the fix: #2097. It is the durable shape rather than the Also carries the counter-contamination half: per-call counters are now The macOS-gated code is typechecked rather than assumed: widening the 15 |
This was referenced Aug 25, 2026
justinchuby
added a commit
that referenced
this pull request
Aug 25, 2026
## Summary - add CUDA `ai.onnx::STFT` v17 using the cuFFT foundation merged in #2080 - preserve the exact CPU/reference and shared shape contracts merged in #2083 - fuse frame extraction, optional windowing, real-to-complex promotion, and batch packing into one NVRTC kernel - execute every `(signal batch, frame)` transform in one cuFFT PlanMany call, then unpack the requested spectrum with one NVRTC kernel ## Claimed surface The CUDA EP claims contiguous f32 signal/output tensors and an optional contiguous f32 window. `frame_step` is a required Int32/Int64 scalar; `frame_length` is an optional Int32/Int64 scalar. Signal shape is `[batch, signal_length, 1|2]`. Real and complex input, full spectrum, and real-input `onesided` (default 1) are supported. Frames are complete, uncentered, and unpadded, exactly matching #2083/ONNX v17. All frames across all signal batches are packed and submitted as one FFT batch. Arbitrary positive frame lengths, including non-power-of-two lengths, are supported by cuFFT. f16/bf16/f64 data, strided signal/window layouts, complex `onesided=1`, malformed scalar ranks/dtypes, zero signal/window extents, and invalid attributes decline at claim time where metadata permits; dynamic zero `frame_step`, mismatched window/frame length, and short signals return typed runtime errors. ## Foundation reuse - Reuses #2080's dynamically loaded cuFFT API, RAII plan ownership, stream binding, external work area, global counters, and bounded 16-entry LRU. DFT and STFT factories receive the same cache instance; no second FFT wrapper/cache exists. - Reuses #2083's already-shared native STFT shape rule and plugin adapter. No shape-rule or adapter census changes were needed. - The plan key includes device, f32/input kind, rank/axis layout, frame length, total FFT batch, and direction. ## Workspace and capture cuFFT auto-allocation remains disabled. Packed frames plus the cuFFT work area come from executor-governed step workspace with overflow-checked offsets and 256-byte alignment. There is no DFT/STFT-side `cudaMalloc`; cuFFT may retain its documented opaque plan/JIT state. CUDA graph capture fails closed with a precise reason: required runtime scalar reads and plan selection are not capture-safe in this implementation. ## GPU validation RTX 4060 Laptop GPU, driver 591.55, pinned CUDA 13.1 wheel environment, serialized: - `cargo test -p onnx-runtime-ep-cuda --features gpu-tests --test stft_gpu -- --test-threads=1` — **9 passed, 0 failed, 0 ignored** - shared cuFFT DFT regression suite — **7 passed, 0 failed, 0 ignored** - CUDA STFT lib/unit selection — **2 passed** - unchanged CPU STFT oracle selection — **13 passed** - shared shape inference STFT selection — **3 passed** - shared plugin shape agreement suite — **4 passed** - source-derived gap guard and covered-op duplicate guard — **1 passed each** - conformance-profile duplicate guard — **1 passed** - `cargo clippy -p onnx-runtime-ep-cuda --all-targets --features gpu-tests -- -D warnings` — passed - `cargo check -p onnx-runtime-ep-cuda-plugin --features cuda` — passed The broad pre-existing `every_covered_op_has_a_conformance_entry` guard remains red only for `PagedAttention`, `Mish`, `Celu`, and `TensorScatter`; the new STFT entry is present and is not in the missing set. ## Mutation evidence Each mutation was applied independently, caught, and reverted: - ignored the window -> `nontrivial_window_is_applied_and_matches_cpu` failed - advanced by `frame_length` instead of `frame_step` -> `overlapping_step_selects_the_middle_frame` failed - omitted the final complete frame -> `final_complete_frame_is_not_dropped` failed - emitted `N/2` instead of `N/2+1` bins -> `onesided_keeps_n_over_two_plus_one_and_matches_full_prefix` failed ## Structural performance/memory evidence For `[batch=2, signal=16, real]`, `frame_length=5`, `frame_step=3`: - frames per signal: **4** - one cuFFT batch: **8 transforms** - explicit custom launches: **2** (fused frame+window pack, output unpack) - cuFFT API executions: **1** (vendor-internal kernel count was not profiled) - governed workspace: **512 bytes** (320 bytes packed complex data, aligned; cuFFT reported zero additional work bytes for this geometry) No latency or throughput claim is made. Remaining optimization opportunities are using the full-spectrum output directly as the in-place FFT buffer, an R2C onesided path, and pre-resolved constant scalars/plans for capture-safe execution. ## Not verified - Linux or H100/H200 execution - CUDA graph capture (explicitly unsupported) - f16/bf16/f64 CUDA STFT - strided CUDA signal/window execution (explicitly declined) - latency/throughput versus CPU, PyTorch, or other GPU FFT stacks - cuFFT vendor-internal kernel launch count Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
justinchuby
added a commit
that referenced
this pull request
Aug 25, 2026
…test (#2093) Closes #2089. Filed by @holden — thanks, the diagnosis in the issue was already correct and I only added the `n >= 4` eligibility detail. ## The bug `kernels::stft::tests::real_unwindowed_overlapping_frames_match_independent_reference` asserts that every power-of-two frame took a fast path, and it does so by reading `DFT_FFT_TEST_HITS`. That counter is bumped **only** by the portable radix-2 branch of `DftPlan::transform`. On Apple targets `DftPlan::new` builds a vDSP setup whenever `n.is_power_of_two() && n >= 4`, and `transform` bumps `DFT_VDSP_TEST_HITS` and **returns** before ever reaching the radix-2 block. The test uses a frame length of 4 — the first vDSP-eligible size — so on macOS the frames took a fast path and the counter the test was watching stayed at zero. Deterministic red on `Rust coverage (macOS arm64)` since 08022b6 (#2083); it is blocking the macOS lane on every open PR, including my own #2072. ## The fix The property the assertion is actually about is *"we did not fall back to the naive O(n^2) transform"*, not *"we took this particular one of the two fast paths"*. So it now reads a `dft_fast_path_hits()` helper that sums both counters. The helper lives in `dft.rs` rather than inline at the call site because `DFT_VDSP_TEST_HITS` is itself `#[cfg(any(target_os = "macos", target_os = "ios"))]` — it does not exist on Linux, so summing it at the `stft.rs` call site would need a second cfg block in a file that has no business knowing about vDSP. **`fft_fallback_reachability` is deliberately untouched.** It uses `n = 2`, which is below the vDSP minimum, so it still exercises and still asserts on the portable radix-2 path specifically, on every target. Weakening that one to use the sum would have thrown away the only real radix-2 check we have. Instead I added `an_eligible_power_of_two_takes_a_fast_path_on_every_target` at `n = 4` — the first vDSP-eligible size — which is the cross-platform counterpart and which nothing previously covered. ## Why this is not a mute button Broadening an assertion to make a red lane green is exactly the shape of change that can silently delete the test, so I did not want to ship it on the argument that it looks right. Neither of us has a Mac, so I reproduced the Apple shape on Linux: I added an early-returning arm for `n >= 4` that bumps a separate counter, mirroring what the vDSP arm does. - **Pre-fix code under the simulation: FAILS**, at `stft.rs:363`, with `each power-of-two frame must use the radix-2 FFT path` — the same line and the same message as the macOS runner. So the simulation is a faithful reproduction, not an approximation. - **Fixed code under the same simulation: PASSES.** - **Radix-2 disabled entirely (`if false && ...`): both `fft_fallback_reachability` and the new test FAIL.** So the fixed assertion still detects a genuinely absent fast path rather than accepting anything. That third arm is the one that matters: it shows the widened assertion has not lost its teeth. A change that only satisfied the first two would be indistinguishable from deleting the check. One more note in the code: the counters are process-global and tests run in parallel, so a concurrent test can inflate `after - before`. The assertion is `>=`, and inflation can only ever mask a *failure to under-count* — which is not the failure mode being guarded — so it cannot produce a false green here. I left a comment saying so rather than reaching for serialisation. ## Validation - `cargo test -p onnx-runtime-ep-cpu --lib` — **1797 passed, 0 failed**, 26 ignored - `cargo test -p onnx-runtime-ep-cpu --lib --features mlas` — **1823 passed, 0 failed**, 37 ignored - `cargo clippy -p onnx-runtime-ep-cpu --lib --all-targets -- -D warnings` — clean - `cargo fmt --all -- --check` — clean Linux can only prove the negative half of this directly; the macOS lane on this PR is the actual verdict, and I will not merge before it is green. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 25, 2026
…2058) (#2098) Closes #2058. 21 crates are compiled and tested by CI and linted by nothing. This derives the clippy package lists from the same source the test lanes already use, and adds a guard so the two cannot drift apart again. ## The correction I owe first #2058 claimed two live `clippy::unnecessary_cast` denials on `main` in `onnx-runtime-ep-plugin/src/compute.rs`. **They are gone, and I did not fix them.** At `bc715c329`: ``` cargo clippy --locked --all-targets -p onnx-runtime-ep-plugin -- -D warnings -> 0 ``` `git log 67d3aa5..HEAD -- crates/onnx-runtime-ep-plugin/` is three feature/refactor commits (#2049, #2064, #2083); none mentions clippy or the cast. The lines were deleted incidentally by a refactor. That weakens the issue's headline and strengthens its actual point. The gap admits defects **and releases them unobserved** — nobody knows what is in there at any moment without running the lint themselves. The sample was never the argument; the mechanism is. ## The rule that created the gap was false `ci.yml` instructed: *"To add a crate, first confirm its normal+dev dependency tree contains no ort-sys/CUDA dependency."* That rule was already violated by the list it annotated. Measured: ``` crates in offline-linux whose normal+dev tree contains onnx-genai-ort-sys: 13 of those, already on the clippy list: onnx-runtime-ep-cpu, onnx-runtime-ep-api ``` And it does not matter which way you resolve that, because **`cargo clippy` only ever checks** — it never links and never runs a test binary. The offline/ort-backed split exists for `cargo test`; it has no force for lint. `onnx-genai-ort-sys` itself compiles here with no network. In #2058 I wrote that I did not know whether the premise was stale or whether `ep-plugin` should be out of `offline-linux`, and would rather flag it than guess. This is the measurement I said I would not substitute a guess for: **the premise was false.** ## What changed - All three clippy `-p` lists — **byte-identical to each other, 31 packages, repeated in three jobs** — become `$(python .github/scripts/workspace_test_packages.py cargo-args lint)`. The new `lint` lane is "every package some test lane compiles" (55). Windows ARM64 takes `offline-linux` (49), still a strict superset of the 31 it linted before, and no ORT crates on that target. - `verify` gains a **lint-coverage half**: it fails if any tested package is reached by no `cargo clippy` invocation anywhere in `.github/workflows`. Generator calls are **expanded, not skipped** — a computed `-p` list counts, and a hand-written one cannot hide behind some other step computing one. - A **self-test step** actually exercises the controls. ## Evidence **The guard detects the real defect.** Run against unmodified `main`, before the `ci.yml` change, it reports the gap by name: ``` Workspace lint coverage check failed. Package(s) are compiled and tested by CI and linted by nothing: - onnx-genai - onnx-runtime-ep-nxrt-abi - onnx-runtime-memory-abi - onnx-genai-capi - onnx-runtime-ep-nxrt-host - onnx-runtime-memory-host - onnx-genai-ort - onnx-runtime-ep-nxrt-testplugin - onnx-runtime-memory-testplugin ... 21 total ``` **The fix closes it, and the pass is not an empty selection.** The exact command `ci.yml` now runs: ``` cargo clippy --locked --all-targets $(python ... cargo-args lint) -- -D warnings -> RC=0 ``` parsed with `--message-format=json`: **wanted 55, seen 55, MISSING: none.** A clippy run that selected nothing would also exit 0, so the package set is confirmed present in the lint graph rather than inferred from the exit code. **Mutations.** | mutation | result | |---|---| | `--simulate-unlinted onnx-runtime-ep-cpu` | rc=1, names that crate | | `--simulate-missing onnx-runtime-ir` (pre-existing half, after refactor) | rc=1 | | both Linux clippy steps reverted to lane `offline-linux` | rc=1, names `onnx-genai`, `onnx-genai-capi`, `onnx-genai-ort` | | a **comment** claiming `cargo clippy -p onnx-genai ...` added, lanes narrowed | rc=1 — prose cannot buy coverage | | clean tree | rc=0 | ## Two defects the mutations found in my own work Recording both, because in each case the check was passing at the time. **1. The scanner counted a YAML comment as an invocation.** After I rewrote the explanatory comment — which contains the words `cargo clippy` — the reported invocation count went `10 -> 11`. Nothing failed; the only symptom was a number moving that I had no reason to expect to move. A comment reading `# cargo clippy -p foo` would have granted `foo` lint coverage. The scanner now skips comment lines, the count is back to 10, and the mutation table above has a cell for exactly this. **2. My first self-test passed because `python` was not on PATH.** It was written as `if cmd ...; then fail; fi` — a bare non-zero check. Command-not-found is `127`, which is non-zero, so it read as *"the guard correctly failed"*. It passed loudest in precisely the case where nothing ran. It now requires **exit 1 specifically, plus the matching failure message**: ``` interpreter present -> rc=0 both guards correctly fail interpreter MISSING -> rc=1 ::error::verify ... exited 127, expected 1 guard neutered -> rc=1 ::error::verify ... exited 0, expected 1 ``` This is the same shape as `QEMU_LD_PREFIX` and `| head`: **a weaker check fails in the direction of passing.** It is also why the self-test step exists at all — `--simulate-missing` had shipped with this guard since it was written and CI had **never once invoked it**. A control nobody runs is not a control. ## Scope and risk - `Fast (Linux x86_64)` and `Rust quality` gain 24 crates each. Verified locally: the whole 55-package set is clean at `-D warnings --all-targets`. - **`Rust (Windows ARM64)` gains 18 crates on `aarch64-pc-windows-msvc`, which I cannot run locally.** If that lane goes red this is mine and I will narrow it before merging. That lane is **not required**, which — as of last night's two merged defects that green required checks did not catch — is exactly the kind of lane I intend to read rather than assume. **I will wait for it.** - No Rust source changes. No behaviour change. Lint coverage only. Normal `--squash --auto`. No admin bypass. --------- 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.
Summary
ai.onnx::STFTon the native CPU EPONNX contract
This implements STFT since_version 17:
signal: T1, required scalarframe_step: T2, optional rank-1window: T1, optional scalarframe_length: T2windoworframe_lengthis required; when both are present their lengths must match, so ONNX STFT has no separate window-length/FFT-length modeonesidedis integer 0/1 and defaults to 1[batch, signal_length, 1|2]; complex input (...x2) requiresonesided=0[batch, floor((signal_length-frame_length)/frame_step)+1, bins, 2], withbins=frame_length/2+1for onesided andbins=frame_lengthotherwiseThe schema permits f16/f32/f64/bf16. This CPU kernel deliberately claims f16, bf16, and f32 only: inputs are widened once, computed in f32, and narrowed once. f64 is rejected at claim and execution rather than pretending f32 arithmetic is f64 support. Length inputs accept int32/int64.
Algorithm and allocation behavior
For
N=frame_length, power-of-two frames use the existing radix-2 FFT (or Accelerate vDSP where available), forO(batch * frames * N log N). Arbitrary lengths currently use the existing scalar DFT,O(batch * frames * N^2); no faster arbitrary-length claim is made.Signal/window strided views are materialized once. The output is allocated once. Four
N-element scratch vectors and one DFT/vDSP plan are reused across all frames and batches; there is no per-frameVecor vDSP-plan allocation.Tests
Mutation checks each went red independently:
windowfailed the strided nontrivial-window oracle (10vs13.75DC)frame_lengthinstead offrame_stepfailed the overlap test+1frame failed the three-frame shape assertionN/2instead ofN/2+1failed the Nyquist/shape testCUDA reuse
A future CUDA STFT can reuse the v17 validation contract, native shape rule, shared plugin adapter, independent f64 reference formula, frame-count/onesided mutation tests, and the fixed-size transform-plan/scratch ownership pattern. This PR does not touch CUDA code or Leon's cuFFT work.
Not verified
The only excluded CPU-lib test is
a_default_width_pool_on_leader_cpus_uses_every_core_it_was_given, which is Linux-affinity-specific and fails on Windows because process-wide affinity masking is intentionally unsupported there.