Fix STFT complex input frame offsets

[GopalakrishnanN]

Description

Fix STFT frame pointer arithmetic for complex-valued input so frame starts are computed in input samples, not trailing real/imag components. Since the frame view pointer is U*, one pointer increment advances one full real or complex sample.

Also add validation that frame_step is positive and keep a defensive bounds check before creating non-owning tensor views.

Review feedback addressed: simplified the frame pointer arithmetic, fixed the swapped STFT input comments, documented the defensive bounds check, and added double-complex regression coverage. The new STFT validation/regression tests exclude kDmlExecutionProvider because these CPU STFT validation/regression paths do not consistently match DirectML behavior in Windows GPU CI.

Motivation and Context

For complex input shaped [batch_size, signal_length, 2], pointer increments already advance by one real/imag pair. Multiplying frame offsets by signal_components == 2 again can advance past the valid frame start, allowing later frames to read across batches or beyond the input allocation.

Testing

• git diff --check -- onnxruntime/core/providers/cpu/signal/dft.cc onnxruntime/test/providers/cpu/signal/signal_ops_test.cc
• .\.venv\Scripts\python.exe tools\ci_build\build.py --config RelWithDebInfo --build --parallel --target onnxruntime_provider_test --build_dir build\Windows
• .\onnxruntime_provider_test.exe --gtest_filter="SignalOpsTest.STFTFloat:SignalOpsTest.STFTFrameStepMustBePositive:SignalOpsTest.STFTFloatComplexInputBatched:SignalOpsTest.STFTDoubleComplexInputBatched" from build\Windows\RelWithDebInfo\RelWithDebInfo

[Copilot]

Pull request overview

Fixes STFT frame pointer arithmetic when the input is complex-valued (shape [batch, signal_length, 2]) so frame offsets are not multiplied by the real/imag component dimension when the underlying pointer is std::complex<T>*, preventing out-of-bounds reads and cross-batch frame reads. Also adds input validation for frame_step and a defensive bounds check before creating non-owning tensor views.

Changes:

• Validate frame_step > 0 and reject invalid inputs early.
• Fix STFT frame start pointer arithmetic for complex input by avoiding double-counting the trailing real/imag dimension.
• Add tests covering invalid frame_step and batched complex input to ensure frames do not bleed across batches.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
onnxruntime/core/providers/cpu/signal/dft.cc	Fixes complex-input STFT frame pointer arithmetic and adds validation/bounds checks.
onnxruntime/test/providers/cpu/signal/signal_ops_test.cc	Adds regression tests for invalid frame_step and batched complex STFT correctness.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[hariharans29]

Verdict: Approve

This is a real OOB read on complex-input STFT and the fix is correct. Test coverage is well-chosen. A couple of small simplification/cleanup observations below, none blocking.

---

Bug analysis — the fix is correct

The dispatch is paired strictly:

short_time_fourier_transform<float,  float>                  // real,    U = float
short_time_fourier_transform<float,  std::complex<float>>    // complex, U = std::complex<float>
short_time_fourier_transform<double, double>                 // real,    U = double
short_time_fourier_transform<double, std::complex<double>>   // complex, U = std::complex<double>

signal_data is U*, so each pointer increment is sizeof(U) bytes — i.e., the trailing real/imag dim is already absorbed into the element type for the complex paths.

Pre-PR arithmetic:

signal_data + (batch_idx * signal_size * signal_components) + (i * frame_step * signal_components);

• Real (U = T, signal_components = 1): correct.
• Complex (U = std::complex<T>, signal_components = 2): every element offset is multiplied by 2 a second time. Byte offset becomes 2 × correct. Cleanly walks off the end of the allocation.

Concretely with the test config (batch=2, signal=128, frame_step=16, frame_length=32, n_dfts=7, signal storage = 256 std::complex<float> elements):

• pre-PR, batch=1, i=6: signal_data + (1*128*2 + 6*16*2) = signal_data + 448 → reads 32 complex elements starting at element 448 of a 256-element allocation. OOB.
• pre-PR, batch=0, i=4: signal_data + (0 + 4*16*2) = signal_data + 128 → reads 32 starting at 128, which is inside batch 1's region. Cross-batch read.

post-PR with signal_frame_step_components = 1:

• batch=1, i=6: signal_data + 128 + 96 = signal_data + 224, reads 224..255. ✓ inside batch 1.
• batch=0, i=4: signal_data + 0 + 64 = signal_data + 64, reads 64..95. ✓ inside batch 0.

Bug and fix verified.

---

Test coverage is solid

STFTFloatComplexInputBatched is well-designed for this exact bug:

• DC-only signal (real lane constant, imag = 0) per batch with different constants per batch (1.0 and 99.0).
• Expected output: DC bin = frame_length × value, all other bins zero. This is the textbook DFT of a constant.
• Cross-batch contamination would jam different constants into adjacent frames and immediately break the DC bin equality — the test cannot pass without the fix. Good targeting.

STFTFrameStepMustBePositive covers both 0 (division-by-zero in n_dfts computation) and -1 (negative stride → backward pointer walks). Both edges, both gated by the new ORT_RETURN_IF_NOT.

---

Observations (non-blocking)

1. signal_frame_step_components is always 1 in the current dispatch

const int64_t signal_frame_step_components = std::is_same<T, U>::value ? signal_components : 1;

• T == U only happens on the real paths, where signal_components == 1 is enforced upstream.
• T != U only happens on the complex paths, where the conditional explicitly returns 1.

So in every reachable call, the value is 1. Two cleaner alternatives:

(a) Just drop the multiplier — signal_data is U* and U already encodes "one sample of whatever the input is":

auto input_frame_begin = signal_data + (batch_idx * signal_size) + frame_start;

(b) Or keep the variable as a documented constant with an assertion:

// signal_data is U*; one increment == one sample, regardless of real vs complex.
constexpr int64_t samples_per_step = 1;
ORT_ENFORCE(std::is_same<T, U>::value ? signal_components == 1 : signal_components == 2,
            "Internal dispatch invariant: U encodes the component dimension.");

The ternary is defensible as defense-in-depth against future dispatch changes (e.g., someone wiring <float, float> to a complex input), but as written it's a guard the reader has to mentally evaluate to discover it's a no-op. If you want to keep it, a one-line comment explaining "redundant in current dispatch; protects against future T==U complex calls" would help the next reader.

2. Pre-existing input-comment error in the same function

Lines 526–530 (just above where you're editing):

// Input(0, "signal") type = T1
// Input(1, "frame_length") type = T2     // <-- actually frame_step
// Input(2, "window") type = T1, optional
// Input(3, "frame_step") type = T2       // <-- actually frame_length

The code below correctly reads ctx->Input<Tensor>(1) as frame_step and ctx->Input<Tensor>(3) as frame_length_tensor, but the header comment has them swapped. Not your bug, but since you're touching the function and adding tests that depend on the actual input ordering, a drive-by comment fix would be a small kindness to future readers. Optional.

3. frame_start <= signal_size - window_size is redundant with n_dfts

The defensive check is fine, but worth noting that:

n_dfts = floor((signal_size - window_size) / frame_step) + 1
max_i = n_dfts - 1
max_frame_start = (n_dfts - 1) * frame_step
                ≤ floor((signal_size - window_size) / frame_step) * frame_step
                ≤ signal_size - window_size

So with frame_step > 0 (now enforced) and window_size <= signal_size (enforced above), the check can never fire. That's fine — defense in depth on a function that produces non-owning tensor views over raw pointers is reasonable — but you may want to add a // defensive; should not fire given n_dfts derivation comment to make the intent explicit, otherwise a reader will spend time chasing the case that triggers it.

4. DML exclusion

The new STFT validation/regression tests exclude kDmlExecutionProvider because these STFT paths are not reliable DirectML assertions in the Windows GPU DML CI job and should be skipped there.

OK as a scope-limiter, but the underlying question is unanswered: does the DML STFT kernel have the same complex-input pointer-arithmetic bug? If yes, this fix only addresses half the problem. If no, then "not reliable" deserves a slightly more specific explanation (does DML fail the validation test because it doesn't validate frame_step? Does it fail the complex-input test because of a separate bug? Or does the test infrastructure flake on DML?). A tracking issue link in the test comment would help. Not blocking this PR, but worth a follow-up.

5. Missing input combinations not covered

The complex-path bug specifically affects std::complex<T>* pointer arithmetic, but the test only covers float complex. The double complex path (short_time_fourier_transform<double, std::complex<double>>) was identically buggy and is identically fixed, but has no regression test. Worth adding a parallel STFTDoubleComplexInputBatched (or templating the test helper) since both code paths share the bug and either could regress independently.

---

Bottom line

Approve. The bug is real, the fix is minimally invasive and correct, and the headline test exercises exactly the cross-batch contamination path. Worth considering: simplify or comment the always-1 ternary (item 1), add a double complex regression test (item 5), and clarify the DML exclusion with a tracking issue (item 4). None of these need to block the merge.

[GopalakrishnanN]

Addressed the non-blocking review feedback in dcdbd19260:

• simplified STFT input frame pointer arithmetic to sample offsets and removed the std::is_same ternary/include
• fixed the swapped STFT input-order comments
• documented why the frame bounds check is defensive
• templated the complex-input regression and added STFTDoubleComplexInputBatched

Validation:

• git diff --check -- onnxruntime/core/providers/cpu/signal/dft.cc onnxruntime/test/providers/cpu/signal/signal_ops_test.cc
• rebuilt onnxruntime_provider_test
• ran SignalOpsTest.STFTFloat, SignalOpsTest.STFTFrameStepMustBePositive, SignalOpsTest.STFTFloatComplexInputBatched, and SignalOpsTest.STFTDoubleComplexInputBatched

[hariharans29]

Verdict: Approve

All five observations from the prior round are addressed. The bug fix is unchanged and still correct; the cleanup makes the code strictly clearer than before; test coverage now includes both float and double complex paths. Ready to merge.

---

Verification of each prior observation

1. Always-1 ternary — fixed cleanly (took the strongest option)

Pre-PR (round 1):

const int64_t signal_frame_step_components = std::is_same<T, U>::value ? signal_components : 1;
auto input_frame_begin = signal_data + (batch_idx * signal_size) + (i * frame_step * signal_frame_step_components);

Post-PR (round 2):

const auto frame_start = i * frame_step;
// signal_data is U*, so one increment advances one input sample, including both lanes for complex input.
auto input_frame_begin = signal_data + (batch_idx * signal_size) + frame_start;

That's option (a) from my prior comment — drop the multiplier and rely on U encoding the per-sample width. Best outcome: less code, no dead ternary, and the inline comment captures the invariant that made the multiplier unnecessary. Also drops the <type_traits> include implicitly.

Sanity-checked the arithmetic against the new test:

• complex, batch=1, i=6: signal_data + 128 + 96 = +224, reads [224..255] of a 256-element std::complex<float> buffer. ✓ in-bounds,