[js/webgpu] Add DFT operator

[nicholascelestin]

Description

Implements the DFT operator (opset 17–20) for the WebGPU (JSEP) execution provider: a batched shared-memory mixed-radix (2/3/4/5) Stockham FFT, with an O(N^2) direct-DFT fallback for non-5-smooth lengths. Handles real/complex input, forward/inverse (1/N-normalized), one-sided RFFT and one-sided-inverse IRFFT, and the opset-20 dft_length/axis tensor inputs — matching the semantics of the CPU kernel in core/providers/cpu/signal/dft.cc.

Motivation and Context

The WebGPU EP had no DFT/FFT kernel, so spectral and audio-preprocessing models fall back to a slow path (#20997). This adds it to the JSEP EP. The op test suite was extended to cover the WebGPU backend; I also verified the kernel against ORT's CPU DFT on real WebGPU across the forward, inverse, RFFT, and IRFFT paths. As with other onnxruntime-web op PRs, CI is the source of truth for the full web build.

Closes #20997

[nicholascelestin]

@microsoft-github-policy-service agree

[nicholascelestin]

@guschmue , know this might take a bit, but was hoping a maintainer could kick off the CI checks, so we've got a place to start?

[nicholascelestin]

@qjia7, just pinging to try and unblock this, since it currently has no reviewers and CI hasn't run yet.

[qjia7]

Thanks for the contribution. Could you also add the WebGPU implementation under the WebGPU provider directory? JSEP will be deprecated soon and replaced by the WASM + WebGPU architecture.

[nicholascelestin]

Thanks for the contribution. Could you also add the WebGPU implementation under the WebGPU provider directory? JSEP will be deprecated soon and replaced by the WASM + WebGPU architecture.

This good?

[qjia7]

Sorry for the late response. I am OOF. Add @hariharans29 to take a look, thanks.

[Copilot]

Pull request overview

Adds ONNX DFT (opset 17–20) support to the WebGPU (JSEP) execution path, enabling FFT/DFT workloads (audio/spectral preprocessing) to run on WebGPU instead of falling back to slower CPU paths.

Changes:

• Registers DFT for the WebGPU EP and implements a mixed-radix Stockham FFT with a direct-DFT fallback.
• Adds corresponding JS/JSEP operator plumbing and a WebGPU TS implementation + test vectors.
• Extends CPU signal op tests and updates WebGPU operator documentation/listing.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
onnxruntime/test/providers/cpu/signal/signal_ops_test.cc	Adds additional DFT test coverage (prime length + zero-padding paths).
onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc	Registers WebGPU DFT kernels for opset 17–20.
onnxruntime/core/providers/webgpu/math/dft.h	Declares WebGPU DFT kernel/program classes.
onnxruntime/core/providers/webgpu/math/dft.cc	Implements WebGPU DFT/FFT shaders and kernel compute logic.
onnxruntime/core/providers/js/operators/dft.h	Adds JS EP kernel attribute wiring for DFT.
onnxruntime/core/providers/js/operators/dft.cc	Registers DFT kernels for the JS EP (opset 17–20).
onnxruntime/core/providers/js/js_execution_provider.cc	Adds DFT to the JS EP kernel registry table.
js/web/test/suite-test-list.jsonc	Adds dft.jsonc to the web test suite list.
js/web/test/data/ops/dft.jsonc	New WebGPU test vectors for DFT forward/inverse/RFFT/IRFFT + opset-20 optional inputs.
js/web/lib/wasm/jsep/webgpu/ops/dft.ts	WebGPU implementation of DFT (FFT + direct fallback) in TS/WGSL.
js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts	Wires DFT into WebGPU op resolve rules.
js/web/docs/webgpu-operators.md	Documents DFT as supported by WebGPU.

Comments suppressed due to low confidence (1)

js/web/lib/wasm/jsep/webgpu/ops/dft.ts:371

• DFT needs to reject non-positive transform lengths. In particular, plan.length === 0 can hang in factorizeToRadices() (since 0 % radix === 0 and remaining /= radix never changes). Adding a simple runtime validation before attempting radix factorization avoids a potential infinite loop and prevents emitting invalid WGSL for negative lengths.

  const plan = computeDftPlan(input, fftAxis, inverse, onesided, dftLength);
  const useSharedMemoryFft = plan.length <= MAX_SHARED_MEMORY_LENGTH && factorizeToRadices(plan.length) !== undefined;
  const programInfo = useSharedMemoryFft ? createFftProgramInfo(plan) : createDirectDftProgramInfo(plan);

[Copilot]

Pull request overview

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

Comments suppressed due to low confidence (2)

onnxruntime/core/providers/webgpu/math/dft.cc:230

• In the IRFFT path, spectrum() assumes that for every k >= uniforms.signal_length, the mirrored bin (N - k) is within the provided half-spectrum. If opset-20 dft_length sets N such that the input half-spectrum is shorter than floor(N/2)+1, length_ - k can be >= uniforms.signal_length and read_sample(length_ - k) will read past the end of the input. CPU DFT treats missing bins as zero-padding. Consider bounding the usable half-spectrum length and returning 0 for bins whose mirror is also out of range.

  if (is_inverse_ && is_onesided_) {
    // For IRFFT the spectrum is the Hermitian extension of the half-spectrum input.
    shader.AdditionalImplementation()
        << "fn spectrum(in_base: u32, k: u32) -> vec2<f32> {\n"
        << "  if (k < uniforms.signal_length) { return " << read_sample("k") << "; }\n"
        << "  let h = " << read_sample(std::to_string(length_) + "u - k") << ";\n"
        << "  return vec2<f32>(h.x, -h.y);\n"
        << "}\n";

js/web/lib/wasm/jsep/webgpu/ops/dft.ts:285

• In the IRFFT direct-DFT path, spectrum() assumes the mirrored bin ${length}u - k always exists in the provided half-spectrum. If opset-20 dft_length sets N such that the input half-spectrum is shorter than floor(N/2)+1, ${length}u - k can be >= uniforms.signalLength and readSample() will index past the end of the input buffer. CPU DFT treats missing bins as zero-padding. Consider clamping the usable half-spectrum length and returning 0 when both k and its mirror are out of range.

    // For IRFFT the spectrum is the Hermitian extension of the half-spectrum input.
    const spectrum =
      inverse && onesided
        ? `fn spectrum(inBase: u32, k: u32) -> vec2<f32> {
    if (k < uniforms.signalLength) { return ${readSample('k')}; }
    let h = ${readSample(`${length}u - k`)};
    return vec2<f32>(h.x, -h.y);
  }`
        : `fn spectrum(inBase: u32, n: u32) -> vec2<f32> {
    if (n < uniforms.signalLength) { return ${readSample('n')}; }
    return vec2<f32>(0.0, 0.0);
  }`;

[hariharans29]

Review: PR #29454 — [js/webgpu] Add DFT operator (head cf66171)

Author: @nicholascelestin. 2 commits: 93aaa69 (initial JSEP-only) → cf66171 (added native WebGPU EP path per @qjia7's request that JSEP is being deprecated). CI: 86 / 86 checks OK. Copilot AI reviewer: 4 comments (1 High, 3 Medium). Reviewers: @edgchen1, @hariharans29, Copilot. Closes issue #20997 (long-standing DFT/STFT WebGPU op request). 12 files, ~1100+ added lines.

Verdict: approve with the High-severity Copilot finding fixed before merge. Substantial, well-designed, correct DFT implementation. The algorithm choice (Stockham mixed-radix FFT with O(N²) direct fallback) is standard and appropriately scoped. Test coverage is thorough. The one blocking issue is a real hang risk from dft_length == 0 triggering an infinite loop in FactorizeToRadices — a one-line fix.

---

What it does

Adds ONNX DFT (opset 17-19 versioned + opset 20 latest) to two paths:

1. Native WebGPU EP (the newer/primary path per @qjia7's guidance):

• onnxruntime/core/providers/webgpu/math/dft.cc (370 lines) — two Program classes (DFTProgram for the FFT path, DFTDirectProgram for the fallback) + a DFT WebGpuKernel dispatcher.
• onnxruntime/core/providers/webgpu/math/dft.h (88 lines).
• onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc — 2 KERNEL_CREATE_INFO entries.

2. JSEP legacy path (kept for backward compatibility):

• js/web/lib/wasm/jsep/webgpu/ops/dft.ts (380 lines) — mirrors the C++ implementation.
• onnxruntime/core/providers/js/operators/dft.{cc,h} + registration in js_execution_provider.cc.
• js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts — TS kernel routing.

3. Tests + docs:

• js/web/test/data/ops/dft.jsonc — 11 JS test vectors (forward/inverse, real/complex, RFFT/IRFFT, opset 17 and 20 with dft_length+axis inputs, negative axis, inner>1 batching, prime-length fallback).
• onnxruntime/test/providers/cpu/signal/signal_ops_test.cc — 8 new C++ CPU-DFT tests (prime length + zero-padded, on both opset 17 and 20).
• js/web/docs/webgpu-operators.md — listing.
• js/web/test/suite-test-list.jsonc — enable dft.jsonc.

---

Algorithm — shared-memory Stockham mixed-radix (2/3/4/5) FFT

Layout:

• One workgroup per (batch × inner) transform row. dispatchGroup.x = batch × inner.
• Workgroup size = 256 threads.
• Shared memory: array<vec2<f32>, 1024> — two 512-slot ping-pong halves.
• Max transform length = 512, must be 5-smooth (composed of 2s, 3s, 4s, 5s) — otherwise falls back to the O(N²) direct DFT below.

Stockham stages (EmitFftStage / stageCode):

• Each stage reads from readOffset and writes to MAX_SHARED_MEMORY_LENGTH - readOffset (ping-pong).
• Twiddle formula: angle = (sign * 2π * j) / (radix * subTransform) * twiddleIndex where twiddleIndex = t % subTransform. Standard Cooley-Tukey/Stockham twiddle indexing.
• Butterfly specializations:
  • Radix-2: out[0] = leg[0] + leg[1]; out[1] = leg[0] - leg[1]. ✓
  • Radix-4: uses a sign-conditional oddRot = (±oddDiff.y, ∓oddDiff.x) (complex multiplication by ±j) to combine two radix-2 stages efficiently. ✓
  • Radix-3, radix-5: direct DFT sum with cos/sin emitted as WGSL literals. ✓
• Output re-indexing: base = (t / subTransform) * (subTransform * radix) + twiddleIndex — the Stockham auto-sort convention that avoids a separate bit-reversal pass. ✓

IRFFT input reconstruction (is_inverse && is_onesided): loads the L=N/2+1 (even N) or L=(N+1)/2 (odd N) input, then mirrors via Hermitian symmetry X[N-k] = conj(X[k]). The conjugate_end is signal_length - 1 for even N (skip the self-conjugate Nyquist bin) or signal_length for odd N. The DC bin (k=0) is never mirrored. ✓

Forward/inverse non-onesided load path: zero-pads if dft_length > signal_length, standard ONNX convention.

Scaling: scale = inverse ? 1/N : 1. Matches ONNX spec (forward unnormalized, inverse 1/N normalized). ✓

---

Fallback — O(N²) direct DFT (DFTDirectProgram)

Used when length is non-5-smooth or > 512. One workgroup per row; each thread computes outputLength / WORKGROUP_SIZE output bins.

Overflow-safe angle indexing (nice detail):

var kn_mod = 0u;
for (var n = 0u; n < N; n++) {
  let angle = -6.283... * f32(kn_mod) / f32(N);  // sign*2π*(k*n mod N)/N
  acc += cmul(spectrum(in_base, n), vec2<f32>(cos(angle), sin(angle)));
  kn_mod += k;
  if (kn_mod >= N) { kn_mod -= N; }
}

The comment "kn_mod tracks (k*n) mod length via addition, so the twiddle index never overflows u32 at large N" explicitly documents the concern. Standard invariant: kn_mod < N and k < N at loop start → after kn_mod += k, kn_mod < 2N → wrap works correctly. ✓

Spectrum function for IRFFT direct fallback: same Hermitian-extension logic as the FFT path. ✓

---

Numerical correctness — verified against test vectors

Manually verified test case "real input, N=5, {1, 2, 3, 4, 5} → DC=15, bin1=(-2.5, 3.44096)":

• DC: 1+2+3+4+5 = 15 ✓
• Bin 1 real: 1 + 2·cos(-72°) + 3·cos(-144°) + 4·cos(-216°) + 5·cos(-288°) = 1 + 0.618 − 2.427 − 3.236 + 1.545 = -2.5 ✓
• Bin 1 imag: 2·sin(-72°) + 3·sin(-144°) + 4·sin(-216°) + 5·sin(-288°) = -1.902 − 1.764 + 2.352 + 4.755 = 3.441 ✓

Test vectors are numerically correct.

Error budget: For length ≤ 512 through log₂(512)=9 stages, error ≈ 9 · f32_epsilon ≈ 1e-6 relative. The tolerance test.SetOutputAbsErr("output", 0.0002f) (2e-4 absolute) is generous but appropriate.

---

Kernel registration — memory placement is correct

ONNX_OPERATOR_VERSIONED_KERNEL_EX(DFT, kOnnxDomain, 17, 19, kWebGpuExecutionProvider,
    (*KernelDefBuilder::Create())
        .TypeConstraint("T1", WebGpuSupportedFloatTypes())    // fp16, fp32
        .TypeConstraint("T2", {int32, int64})                 // dft_length dtype
        .InputMemoryType(OrtMemTypeCPU, 1),                   // dft_length → CPU
    DFT);

ONNX_OPERATOR_KERNEL_EX(DFT, kOnnxDomain, 20, kWebGpuExecutionProvider,
    (*KernelDefBuilder::Create())
        // ... same as above ...
        .InputMemoryType(OrtMemTypeCPU, 1)   // dft_length → CPU
        .InputMemoryType(OrtMemTypeCPU, 2),  // axis → CPU (new in opset 20)
    DFT);

Both scalar-tensor inputs (dft_length, axis) are marked OrtMemTypeCPU because they're read on the host via TryReadScalar. Correct — no useless GPU roundtrip.

---

🔴 Copilot comment #2 (High) — real bug, please fix before merge

dft_length should be validated as a scalar and required to be > 0 (CPU kernel errors on <= 0). Allowing 0 can also lead to undefined behavior in the 5-smooth factorization path (e.g., FactorizeToRadices(0) loops forever).

Verified by inspection of dft.cc:

static bool FactorizeToRadices(uint32_t length, std::vector<uint32_t>& radices) {
  uint32_t remaining = length;
  for (uint32_t radix : {4u, 2u, 3u, 5u}) {
    while (remaining % radix == 0) {   // for length=0: 0%4==0 always true
      radices.push_back(radix);
      remaining /= radix;              // 0/4 == 0, stays at 0
    }
  }
  return remaining == 1;
}

With length=0, the inner while (remaining % radix == 0) never exits — 0 % 4 == 0 is always true, and 0 / 4 == 0 keeps remaining at 0. The loop pushes 4u into radices indefinitely until the vector allocation OOMs or the runtime abort/timeout kicks in. Effectively a hang on the CPU shader-generation path.

The current guard misses this:

ORT_RETURN_IF(length < 0 || length > std::numeric_limits<uint32_t>::max(), "Invalid DFT length: ", length);

length == 0 bypasses this. Copilot's suggested fix (length <= 0) is correct and one-character:

ORT_RETURN_IF(length <= 0 || length > std::numeric_limits<uint32_t>::max(), "Invalid DFT length: ", length);

This is the only blocking item. Please apply before merge.

---

Copilot comments #1, #3, #4 (Medium, defensive) — nice-to-have

• Set up CI with Azure Pipelines #1: Validate that axis (opset 20 input 2) is a scalar tensor before reading. Current code silently reads element 0 regardless of shape. ONNX spec constrains it to scalar so this is a defense-in-depth ask. Non-blocking.
• update HighLevelDesign.md #3: readScalar() in dft.ts uses Number(getBigInt64Array()[0]) without safe-integer check. For dft_length > 2^53 the value loses precision. Theoretical — real DFT lengths are far smaller. Non-blocking.
• Incremental updates. #4: TS validateInputs() should check dims.length >= 2 — the C++ side does (ORT_RETURN_IF(rank < 2, ...)), but the TS side would crash accessing dims[dims.length - 1] on a rank-0/rank-1 tensor. Cheap fix, matches C++ symmetry. Non-blocking.

Would be nice to apply all three in a follow-up. Only #2 blocks.

---

Test coverage — thorough

11 JS test vectors in dft.jsonc cover:

• Forward complex-complex N=4 (basic sanity).
• Forward real input N=5 (real → complex path).
• Batched axis=1 of [2, 3, 2] (batch × transform).
• Prime length N=7 (direct DFT fallback path).
• inner > 1: middle axis of [1, 3, 2, 2] (2-D DFT column pass — exercises the inner broadcast between axis and innermost complex dim).
• Inverse complex→complex (1/N normalization).
• One-sided forward / RFFT (N=6 → 4 bins).
• One-sided inverse / IRFFT (4 half-bins → 6 real).
• Opset 20 with dft_length and axis tensor inputs (zero-padded N=5→8).
• Negative axis attribute (opset 17).
• Opset 20 with both optional inputs omitted (spec default axis=-2).

8 C++ CPU tests in signal_ops_test.cc:

• DFT{17,20}_Float_prime_length[_onesided] — 4 tests exercising the direct-DFT fallback with prime length 7.
• DFT{17,20}_Float_zero_padded_{radix2,prime} — 4 tests exercising dft_length zero-padding to both radix-friendly (8) and non-radix (7) targets.

The CPU tests will be picked up by the WebGPU EP on the WebGPU CI leg — implicit cross-check. Absolute error tolerance 2e-4 is appropriate.

Coverage gap (non-blocking): No fp16 test. WebGpuSupportedFloatTypes() enrolls fp16, and the shader casts through vec2<f32> internally, so fp16 should work — but it's unexercised. Worth a follow-up test.

---

Minor observations (non-blocking)

1. JSEP and native WebGPU implementations are near-identical. Same algorithm, same shared-memory layout (2×512 vec2), same workgroup size (256). Good discipline for cross-path consistency.
2. Shader cache hint in dft.cc includes (length, input_components, output_components, is_inverse, is_onesided) — all the fields that affect emitted WGSL. batch/signal_length/inner/output_length are passed as uniforms and don't affect the cache key. Program name splits FFT vs direct ("DFT" vs "DFTDirect"). Correct. ✓
3. WgslFloat precision — 9 significant digits (both JS toPrecision(9) and C++ std::setprecision(9)). Matches f32's ~7 sig digits + margin. ✓
4. The workgroup_idx/workgroup_index shader placeholder differs between C++ (workgroup_idx) and JS (workgroup_index). Both substituted by the respective ORT shader helper to workgroup_id.x. Consistent within each path.
5. .SetWorkgroupSize(256) is the hardcoded workgroup size for both programs. For very small outputLength (say 7 in the prime-length test), most of the 256 threads are idle in the outer for (var k = local_idx; k < outputLength; ...) loop. Not optimal but not wrong. Could be tuned in a follow-up.
6. Comment // Bluestein chirp-z ... would restore O(N log N) if the fallback ever becomes hot" — good forward-looking note in the header comment. Signals the author knows the direct DFT is a stopgap and where the follow-up work would go.
7. Long-standing issue closed: [Feature Request] DFT/STFT WebGPU op support (web/js) #20997 has been open for months. This is a substantive contribution.
8. CI 86/86 OK — passes on the entire ORT matrix at the current head.

---

Bottom line

Well-executed DFT implementation that closes a real gap in the WebGPU EP. Algorithm is standard (Stockham mixed-radix FFT + O(N²) fallback), shader is efficient (shared-memory ping-pong, per-radix butterfly specializations), test coverage is thorough across forward/inverse/RFFT/IRFFT/opset-17-vs-20/prime-length/zero-padded/batched-inner variants, and the JS+C++ implementations are structurally near-identical to reduce cross-path drift. The author responded correctly to @qjia7's feedback about JSEP deprecation by adding the native WebGPU EP path alongside the initial JSEP-only implementation.

Ask before merge: apply Copilot comment #2 (change length < 0 → length <= 0 in the ORT_RETURN_IF guard) to prevent the FactorizeToRadices(0) infinite loop. Ideally also apply #1, #3, #4 for defense-in-depth, but only #2 is a hard bug.

[hariharans29]

Could you please manually resolve already addressed copilot comments ? Thanks

[nicholascelestin]

Could you please manually resolve already addressed copilot comments ? Thanks

Should be done!

[hariharans29]

Could you please manually resolve already addressed copilot comments ? Thanks

Should be done!

Thanks - it was resolved when I refreshed my page. I think there are 2 more copilot comments and some minor observations in my review comment. Can you please check if they are relevant / adressable ?

[nicholascelestin]

Thanks - it was resolved when I refreshed my page. I think there are 2 more copilot comments and some minor observations in my review comment. Can you please check if they are relevant / adressable ?

Copilot comments #1-#4 on the validation issues are addressed.
Copilot comments #5-6 on the IRFFT potential mismatch between assumed length and chosen length are now addressed with tests. Issue is actually present in CPU kernel, so new radix-2 length test cases exempts CPU EP. Happy to file issue, if that would help.
Added fp16 tests per your comment.
Would like to defer Bluestein for now, as well as dynamic workgroups (like the top_k op) for small output lengths, but willing to add either if required to merge.
Think the Windows CI failures are unrelated — DFT tests passed on those windows jobs.

[nicholascelestin]

CI failures my fault this time. CPU-excluded radix-2 test cases now skipped explicitly with GTEST_SKIP if no eligible EP (now just WebGPU).