[Performance] DFT/STFT CPU performance is 100-700x slower than PyTorch FFT for batched short transforms

### Describe the issue

I was debugging why ONNX was performing very poorly on a relatively small CNN model and ended up narrowing it down to STFT/DFT operators. My deep model uses STFT to convert a batched time-series to spectrograms at the input side. Benchmarks indicated that STFT/DFT operators in ONNX are significantly slow.

I made a simple repro code (attached) where I just have a PyTorch model with a single computation: STFT. I compared 3 cases: base STFT, a DFT with an unfolded input, and a convolution based DFT implementation. I wanted to narrow down any potential export bugs or something.

The convolution based input is the only one that is not 700x slower (though still is about 5x slower). The other two are orders of magnitude worse. This is extremely bad when the DFT parameters are not powers of 2 (in example, when length = 1000, nfft = 100, hop = 10). PyTorch doesn't suffer from the same performance degradation regardless of the parameters.

### To reproduce

Here is my pyproject:

```toml
[project]
name = "onnx-bug-repro"
version = "0.1.0"
requires-python = ">=3.14"
dependencies = [
"numpy>=2.5.2",
"onnx>=1.22.0",
"onnxruntime-gpu>=1.29.0",
"onnxscript>=0.7.1",
"torch~=2.13.0",
"torchvision~=0.28.0",
]

[[tool.uv.index]]
name = "pytorch-cu132"
url = "https://download.pytorch.org/whl/cu132"
explicit = true

[tool.uv.sources]
torch = [
{ index = "pytorch-cu132", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
torchvision = [
{ index = "pytorch-cu132", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
]
```

```python

import math
from time import perf_counter

import onnx
import onnxruntime as ort
import torch
import torch.nn.functional as F
from torch import nn

class STFT(torch.nn.Module):
def __init__(self, nfft: int, hop: int) -> None:
super().__init__()
self.window = nn.Buffer(torch.hann_window(nfft), persistent=False)
self.nfft = nfft
self.hop = hop

def forward(self, x: torch.Tensor) -> torch.Tensor:
return torch.stft(
x,
n_fft=self.nfft,
hop_length=self.hop,
window=self.window,
center=False,
return_complex=True,
)

class FramedRFFT(torch.nn.Module):
def __init__(self, nfft: int, hop: int) -> None:
super().__init__()
self.window = nn.Buffer(
torch.hann_window(nfft),
persistent=False,
)
self.nfft = nfft
self.hop = hop

def forward(self, x: torch.Tensor) -> torch.Tensor:
frames = x.unfold(
dimension=-1,
size=self.nfft,
step=self.hop,
)

frames = frames * self.window

return torch.fft.rfft(frames, dim=-1)

class ConvSTFT(torch.nn.Module):
def __init__(self, nfft: int, hop: int) -> None:
super().__init__()

n_bins = nfft // 2 + 1

window = torch.hann_window(nfft)

# Build the Fourier basis in float64, then cast the finished kernel to
# float32. Computing the phase directly in float32 introduces enough
# sin/cos error at the higher bins to make this implementation differ
# from torch.stft by more than 1e-4.
frequency = torch.arange(n_bins, dtype=torch.float64)[:, None]
time = torch.arange(nfft, dtype=torch.float64)[None, :]

phase = 2.0 * math.pi * frequency * time / nfft

real_kernel = (torch.cos(phase) * window.double()).float()
imag_kernel = (-torch.sin(phase) * window.double()).float()

# [2 * n_bins, 1, nfft]
kernel = torch.cat(
(real_kernel, imag_kernel),
dim=0,
).unsqueeze(1)

self.kernel = nn.Buffer(kernel, persistent=False)
self.n_bins = n_bins
self.hop = hop

def forward(self, x: torch.Tensor) -> torch.Tensor:
# x: [batch, samples]
spectrum = F.conv1d(
x.unsqueeze(1),
self.kernel,
stride=self.hop,
)

real = spectrum[:, : self.n_bins]
imag = spectrum[:, self.n_bins :]

# Same float representation ONNX uses for complex values:
# [batch, bins, frames, 2]
return torch.stack((real, imag), dim=-1)

def main():
nfft: int = 100 # Set to 128 for power-of-2
hop: int = 10 # Set to 8 for power-of-2

warmup: int = 50
num_runs: int = 500

device = torch.device("cpu")

stft = STFT(nfft, hop).eval().to(device)

batch_size: int = 100
signal_length: int = 1000 # Set to 1024 for power-of-2

# 1. Basic STFT operator

example_input = torch.randn(batch_size, signal_length, device=device)

for _ in range(warmup):
with torch.inference_mode():
_ = stft(example_input)

start = perf_counter()
for _ in range(num_runs):
with torch.inference_mode():
_ = stft(example_input)
end = perf_counter()

print(
f"Average time per run for base torch: {(end - start) / num_runs:.6f} seconds"
)

onnx_stft = torch.onnx.export(stft, (example_input,), dynamo=True)

if onnx_stft is None:
raise RuntimeError("ONNX export failed")

model_proto = onnx_stft.model_proto

onnx.checker.check_model(model_proto)

print(onnx.printer.to_text(model_proto.graph))

options = ort.SessionOptions()
options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
ort_session = ort.InferenceSession(
model_proto.SerializeToString(), options, providers=["CPUExecutionProvider"]
)

for _ in range(warmup):
ort_session.run(None, {"x": example_input.numpy()})

ort_start = perf_counter()

for _ in range(num_runs):
ort_session.run(None, {"x": example_input.numpy()})
ort_end = perf_counter()

print(
f"Average time per run for ONNX: {(ort_end - ort_start) / num_runs:.6f} seconds"
)

# 2. Framed DFT based approach

framed_stft = FramedRFFT(nfft, hop).eval().to(device)

for _ in range(warmup):
with torch.inference_mode():
_ = framed_stft(example_input)

start = perf_counter()
for _ in range(num_runs):
with torch.inference_mode():
_ = framed_stft(example_input)
end = perf_counter()

print(
f"Average time per run for base torch (framed rfft):"
f" {(end - start) / num_runs:.6f} "
f"seconds"
)

program = torch.onnx.export(
framed_stft,
(example_input,),
dynamo=True,
opset_version=20,
)

if program is None:
raise RuntimeError("Failed to export model")

model_proto = program.model_proto

onnx.checker.check_model(model_proto)

print(onnx.printer.to_text(model_proto.graph))

options = ort.SessionOptions()
options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
ort_session = ort.InferenceSession(
model_proto.SerializeToString(),
options,
providers=["CPUExecutionProvider"],
)

for _ in range(warmup):
ort_session.run(None, {"x": example_input.numpy()})

ort_start = perf_counter()

for _ in range(num_runs):
ort_session.run(None, {"x": example_input.numpy()})
ort_end = perf_counter()

print(
f"Average time per run for ONNX (framed rfft): {(ort_end - ort_start) / num_runs:.6f} "
f"seconds"
)

# 3. Convolution based approach

conv_stft = ConvSTFT(nfft, hop).eval().to(device)

for _ in range(warmup):
with torch.inference_mode():
_ = conv_stft(example_input)

start = perf_counter()
for _ in range(num_runs):
with torch.inference_mode():
_ = conv_stft(example_input)
end = perf_counter()

print(
f"Average time per run for base torch (convolution):"
f" {(end - start) / num_runs:.6f} "
f"seconds"
)

program = torch.onnx.export(
conv_stft,
(example_input,),
dynamo=True,
opset_version=20,
)

if program is None:
raise RuntimeError("Failed to export model")

model_proto = program.model_proto

onnx.checker.check_model(model_proto)

print(onnx.printer.to_text(model_proto.graph))

ort_session = ort.InferenceSession(
model_proto.SerializeToString(),
options,
providers=["CPUExecutionProvider"],
)

for _ in range(warmup):
ort_session.run(None, {"x": example_input.numpy()})

ort_start = perf_counter()

for _ in range(num_runs):
ort_session.run(None, {"x": example_input.numpy()})
ort_end = perf_counter()

print(
f"Average time per run for ONNX (convolution): {(ort_end - ort_start) / num_runs:.6f} "
f"seconds"
)

expected = torch.view_as_real(stft(example_input))
actual = conv_stft(example_input)

torch.testing.assert_close(
actual,
expected,
rtol=1e-4,
atol=1e-4,
)

if __name__ == "__main__":
main()

```

Output will be something like:
For powers of 2 (nfft = 128, hop = 8, signal_length = 1024):
```plaintext
Average time per run for base torch: 0.000285 seconds
Average time per run for ONNX: 0.045795 seconds
Average time per run for base torch (framed rfft): 0.000268 seconds
Average time per run for ONNX (framed rfft): 0.053239 seconds
Average time per run for base torch (convolution): 0.000759 seconds
Average time per run for ONNX (convolution): 0.003028 seconds
```

So the base STFT and DFT approaches here are ~200x slower where as the conv method is ~4x slower.

For non-powers of 2 (nfft = 100, hop = 10, signal_length = 1000):

```plaintext
Average time per run for base torch: 0.000217 seconds
Average time per run for ONNX: 0.162501 seconds
Average time per run for base torch (framed rfft): 0.000273 seconds
Average time per run for ONNX (framed rfft): 0.167483 seconds
Average time per run for base torch (convolution): 0.000448 seconds
Average time per run for ONNX (convolution): 0.002244 seconds
```

So the base STFT and DFT approaches here are ~750x slower where as the conv method is ~5x slower.

### Urgency

_No response_

### Platform

Windows

### OS Version

Windows 11

### ONNX Runtime Installation

Released Package

### ONNX Runtime Version or Commit ID

onnxruntime-gpu==1.29.0

### ONNX Runtime API

Python

### Architecture

X64

### Execution Provider

Default CPU

### Execution Provider Library Version

NA

### Model File

Model is very simple and generated as part of the script above.

### Is this a quantized model?

No

Start by running the supplied Python repro with the CPUExecutionProvider and compare the generated graphs for torch.stft, framed torch.fft.rfft, and the convolution implementation. Profile the slow operators for power-of-two and non-power-of-two transforms, then verify that the CPU path approaches the convolution baseline without changing outputs or breaking the ONNX model checks.

Written by the indexing model from the issue text.