[Theodore Cooper (the0cp)]

Fix #32403.

Description

The CPU STFT kernel currently computes each frame independently. This is slow for batched short transforms, especially for non-power-of-two DFT sizes where the DFT path uses Bluestein internally.

We already have an STFTDecomposition transformer that can express real-valued STFT as standard ONNX ops, but it was not registered in the default CPU graph optimization path. This PR makes that decomposition available at Level2 with a CPU EP filter, so it runs after EP partitioning and only rewrites STFT nodes assigned or fallbacked to CPU.

The rewrite is:

STFT
=> Reshape -> Conv -> Reshape -> Transpose

For constant float windows, the window is pre-multiplied into the generated Conv weights. Large generated Conv weights are capped at 64 MiB so large transforms remain as native STFT nodes instead of materializing an O(N^2) initializer.

Local benchmark

Environment:

Fedora 44
CPU: 11th Gen Intel(R) Core(TM) i5-11260H @ 2.60GHz
Build: Debug
Provider: CPUExecutionProvider
Threads: intra-op=1, inter-op=1

Benchmark compares the same STFT model with graph optimizations disabled versus STFT decomposition enabled.

Model	Optimization	Avg latency	P50 latency	Throughput
nfft=100, hop=10, signal_length=1000, batch=100	disabled	2492.61 ms	2420.96 ms	0.40/s
nfft=100, hop=10, signal_length=1000, batch=100	STFTDecomposition	29.36 ms	28.41 ms	34.05/s
nfft=128, hop=8, signal_length=1024, batch=100	disabled	574.96 ms	570.40 ms	1.74/s
nfft=128, hop=8, signal_length=1024, batch=100	STFTDecomposition	48.00 ms	46.38 ms	20.83/s
nfft=1024, hop=256, signal_length=8192, batch=64	disabled (-o 0)	1093.18 ms	1076.51 ms	0.91/s
nfft=1024, hop=256, signal_length=8192, batch=64	Level2 (-o 2)	94.27 ms	93.84 ms	10.61/s
nfft=2048, hop=512, signal_length=16384, batch=64	disabled (-o 0)	2327.06 ms	2307.76 ms	0.43/s
nfft=2048, hop=512, signal_length=16384, batch=64	Level2 (-o 2)	304.03 ms	298.09 ms	3.29/s

Observed local speedups:

• nfft=100: about 84.9x
• nfft=128: about 12.0x
• nfft=1024: about 11.6x
• nfft=2048: about 7.7x

The nfft=1024/2048 data points show that the decomposition still helps beyond very small transforms, while the generated-weight cap keeps larger O(N^2) cases on the native STFT path.

The optimized graph contains:

{'Reshape': 2, 'Conv': 1, 'Transpose': 1}

instead of retaining an STFT node.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

🟡 Changes recommended

Critical correctness and resource-usage issues, plus execution-provider gating, must be addressed before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Enables default Level1 STFT decomposition to improve CPU performance using convolution-based graphs.

Changes:

• Registers STFTDecomposition at Level1.
• Expands decomposition support and optimization.
• Adds transformation and numerical-equivalence tests.

File summaries

File	Review
onnxruntime/test/optimizer/graph_transform_test.cc	Adds STFT decomposition tests; a two-sided (onesided=0) numerical case is still needed (nit).
onnxruntime/core/optimizer/stft_decomposition.cc	Critical: mishandles int32 initializers, does not validate the ONNX domain/schema, and lacks a size limit for quadratic Fourier weights.
onnxruntime/core/optimizer/graph_transformer_utils.cc	Moderate: unrestricted registration may replace native STFT nodes on DML and introduce cross-device transfers.

Review details

Suppressed comments (2)

onnxruntime/core/optimizer/graph_transformer_utils.cc:283

• Registering this transformer at Level1 duplicates the existing STFTDecomposition(dml_ep) registration for DML MCDM devices in core/session/inference_session.cc:2663. GraphTransformerManager::Register rejects duplicate names at the same level (graph_transformer_mgr.cc:87-95), so session initialization fails with “already registered” whenever Level1 optimizations are enabled on those devices. Consolidate the registrations while preserving the DML fp16-specific path.

      transformers.emplace_back(std::make_unique<STFTDecomposition>());

onnxruntime/core/optimizer/stft_decomposition.cc:256

• The rewritten channel packing also handles onesided=0, but every added numerical test explicitly uses onesided=1. Add a two-sided case so the new {batch, 2, dft_size, frames} reshape and {0, 3, 2, 1} transpose are verified for the full-spectrum output.

    const int64_t output_num_frames = ((signal_length - dft_size) / frame_step_value) + 1;
    auto dft_unique_bins = is_onesided ? ((dft_size >> 1) + 1) : dft_size;

• Files reviewed: 3/3 changed files
• Comments generated: 4
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

The unresolved semantic mismatch and provider eligibility regression must be fixed before approval.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 2
• Review effort level: Balanced

[Tianlei Wu (tianleiwu)]

Thanks for tackling this - the 85x number is real and the rewrite is a genuine improvement over the old Conv+Conv+Concat shape.

I verified the new decomposition math independently and it is correct:

• Output channel k uses cos(-2*pi*k*n/N)*w[n] -> Re{sum x*e^{-i2*pi*k*n/N}}, channel bins+k uses sin(...) -> Im{...}.
• Reshape [B, 2*bins, 1, F] -> [B, 2, bins, F] preserves the flat layout because the channel axis splits exactly as c = i*bins + k.
• Transpose perm={0,3,2,1} yields [B, F, bins, 2], matching the ONNX STFT-17 output layout.
• output_num_frames = (signal_length - dft_size)/frame_step + 1 matches both ONNX shape inference and the CPU kernel's n_dfts (dft.cc:585).

Deriving output_num_frames / dft_unique_bins from the inputs instead of reading the STFT node's output shape is a solid robustness win, and the new null/range guards remove several unconditional Shape() dereferences. The registration gate is also mechanically correct: the CPU EP is appended at inference_session.cc:2455, before AddPredefinedTransformers at :2608, so NumProviders() == 1 && Get(kCpuExecutionProvider) really does mean CPU-only and the DML MCDM registration at :2663 is not double-registered.

Three things I'd like resolved before merge, detailed inline:

1. Weight precision. The float phase accumulation in the weight loop produces up to 6.2e-4 absolute error per weight at dft_size=2048 (1.7e-4 at Whisper's 400). Latent before, but this PR turns it on by default for every CPU session. One-line fix.
2. Level1 vs Level2. This is the first Level1 entry whose presence depends on the registered EP set, a few lines below a comment stating Level1 does no EP filtering. Level2 + cpu_ep would also cover the mixed-EP sessions that are currently excluded (a CUDA/QNN session's STFT node partitions to CPU anyway, so it keeps the slow path), and would keep the decomposition out of optimized_model_filepath / ORT-format output.
3. clang-format. The new registration block is not clang-format clean and will fail the lint check.

Test coverage gaps

The breadth here is good (3-D/2-D signal, constant vs runtime window, int32 scalars, onesided=0, double skip, the size cap, the DML EP-set split, and a RunsAtLevel1ByDefault test that exercises the registration rather than the transformer in isolation), and TransformerTester checks numerics against the untransformed baseline rather than just op counts. Missing:

• Any case with dft_size >= 256. Every numeric test uses dft_size in {5, 10} with a 3e-4 tolerance, which is precisely why the precision issue above is invisible. A dft_size=512 case with a tight tolerance would lock the fix in.
• The dft_size > signal_length skip added in this PR.
• A non-constant frame_step (should skip), and a runtime window whose declared length disagrees with frame_length.

Benchmark

The numbers cover nfft = 100 and 128 only. A data point at nfft = 1024/2048 would show where the O(N^2) Conv stops beating the kernel's radix/Bluestein FFT, and would justify where the size cap is set.

[Theodore Cooper (the0cp)]

Thanks for the detailed review. I addressed the concrete issues in the latest commit:

• moved the default registration from Level1 CPU-only gating to Level2 with cpu_ep
• changed phase generation to use double and modulo reduction
• added the static runtime-window length mismatch guard
• removed the unused include
• added tests for dft_size=512, dft_size > signal_length, and non-constant frame_step
• added local nfft=1024/2048 benchmark data to the PR description.

For the 64 MiB cap, I kept the current conservative cap for now and added the requested benchmark data.

[Copilot]

🔵 Needs a closer look

Valid three-input window-only STFT nodes are incorrectly excluded from the optimization.

Review details

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

onnxruntime/core/optimizer/stft_decomposition.cc:197

• A valid window-only STFT has exactly three inputs (signal, frame_step, window), so this guard always leaves that form on the slow native CPU path. The repository already constructs this standard form in onnxruntime/test/providers/qnn/stft_op_test.cc:72-83, and the code below is intended to derive dft_size from the window when frame_length is absent. Accept three inputs and represent a missing fourth input safely before looking up its initializer; add a window-only Level2 decomposition case.

• Files reviewed: 3/3 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[Tianlei Wu (tianleiwu)]

The follow-up commits address the previous precision, dynamic window/frame-length, initializer-size, optimization-level, formatting, and edge-coverage concerns. The Level2 CPU-filtered rewrite preserves STFT-17 layout and scalar semantics, and all reported CI checks pass.

[Theodore Cooper (the0cp)]

The CI failures appear unrelated to this PR, could someone with permission please rerun the failed jobs?

[Theodore Cooper (the0cp)]

Tianlei Wu (@tianleiwu) The merge conflict is resolved by keeping both includes, with no changes to the PR logic, could you please re-approve and re-enable auto-merge when convenient? Thanks!