sync 91d990483a17

Browse files

Files changed (13) hide show

README.md +15 -11
build/webgpu/bench.json +0 -1
build/webgpu/dft-tiled-real.wgsl.jinja +9 -10
build/webgpu/fft-radix2-dit-shared-inplace.wgsl.jinja +15 -16
build/webgpu/fft-radix2-dit-storage-pass.wgsl.jinja +35 -38
build/webgpu/fft-radix2-dit-storage.wgsl.jinja +8 -8
build/webgpu/fft-real-packed-stockham-shared.wgsl.jinja +2 -2
build/webgpu/fft-stockham-shared.wgsl.jinja +32 -44
build/webgpu/fft-stockham-twiddles.wgsl.jinja +9 -10
build/webgpu/manifest.json +199 -415
build/webgpu/metadata.json +29 -14
build/webgpu/stft-general.wgsl.jinja +22 -23
build/webgpu/test.json +6 -7

README.md CHANGED Viewed

@@ -18,20 +18,20 @@ See the [ONNX `STFT` spec](https://onnx.ai/onnx/operators/onnx__STFT.html) for t
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
-| --- | --- | --- | --- | --- | --- | --- |
-| `signal` | `signal` | `T` | `3` | — | Real input signal with shape `(batch_size, signal_length, 1)`, or complex input with shape `(batch_size, signal_length, 2)`. | required |
-| `window` | `window` | `T` | `1` | — | Optional 1-D window tensor with shape `(window_length)` to multiply by each frame before the DFT. | optional |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
-| --- | --- | --- | --- | --- | --- | --- |
-| `output` | `output` | `T` | `4` | derived; see description | STFT result with shape `(batch_size, frames, dft_unique_bins, 2)` in one-sided mode or `(batch_size, frames, frame_length, 2)` in two-sided mode, with real and imaginary parts in the last dimension. `frames = floor((signal_length - frame_length) / frame_step) + 1`; partial trailing frames are not emitted. | required |
 ## Runtime arguments
-| Name | Kind | Semantic | Description | Presence |
 | --- | --- | --- | --- | --- |
 | `frameStep` | `i32` | `frame_step` | Positive number of signal positions between the starts of consecutive frames. | required |
 | `frameLength` | `i32` | `frame_length` | Optional number of signal values in each frame; must equal the window length when both are supplied. | optional |
@@ -52,7 +52,7 @@ Default values (overridable per request):
 ## Files
-- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
@@ -67,10 +67,14 @@ Default values (overridable per request):
 ## Use with `@huggingface/kernels`
-The loader derives every required output's shape and logical dtype from the manifest contract and this call.
-It then allocates the result tensors automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

 ## Inputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `signal` | `T` | `3` | — | Real input signal with shape `(batch_size, signal_length, 1)`, or complex input with shape `(batch_size, signal_length, 2)`. | required |
+| `window` | `T` | `1` | — | Optional 1-D window tensor with shape `(window_length)` to multiply by each frame before the DFT. | optional |
 ## Outputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `output` | `T` | `4` | derived | STFT result with shape `(batch_size, frames, dft_unique_bins, 2)` in one-sided mode or `(batch_size, frames, frame_length, 2)` in two-sided mode, with real and imaginary parts in the last dimension. `frames = floor((signal_length - frame_length) / frame_step) + 1`; partial trailing frames are not emitted. | required |
 ## Runtime arguments
+| Name | Kind | Upstream attribute | Description | Presence |
 | --- | --- | --- | --- | --- |
 | `frameStep` | `i32` | `frame_step` | Positive number of signal positions between the starts of consecutive frames. | required |
 | `frameLength` | `i32` | `frame_length` | Optional number of signal values in each frame; must equal the window length when both are supplied. | optional |
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.STFT",
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
   "cases": [
     {

 {
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
   "cases": [
     {

build/webgpu/dft-tiled-real.wgsl.jinja CHANGED Viewed

@@ -5,18 +5,17 @@
 // phase drift without transcendental functions in the inner loop.
 {{ env.wgsl.resourceDeclarations }}
 // Backend-stable unit-circle twiddles for DFT kernels. Retain the phase as a
-// rational turn and route all polynomial evaluation through the shared precise
-// trig core so FFT, direct DFT, STFT, and cosine windows agree exactly.
 //
 // phase is in [0, period). sign is -1 for a forward transform and +1 for an
 // inverse transform. The result is (cos, sin).
 // Backend-stable f32 sine/cosine core.
 //
-// Shader transcendental accuracy is implementation-defined, and some portable
-// backends are only accurate to roughly 1e-4. Each path retains the most accurate
-// available phase representation, reduces it to [-pi, pi], then uses these
-// polynomials. The half-pi core is degree 13 for sine and degree 12 for cosine;
-// truncation error is well below one f32 ULP over its documented interval.
 const PRECISE_TRIG_PI: f32 = 3.141592653589793;
 const PRECISE_TRIG_TWO_PI: f32 = 6.283185307179586;
 const PRECISE_TRIG_HALF_PI: f32 = 1.5707963267948966;
@@ -84,9 +83,9 @@ fn dft_precise_twiddle(phase: u32, period: u32, sign: f32) -> vec2<f32> {
   return precise_sincos_centered(sign * PRECISE_TRIG_TWO_PI * centeredTurns);
 }
-const N: u32 = {{ source.transformLength }}u;
-const NUM_BINS: u32 = {{ source.numBins }}u;
-const WG: u32 = {{ source.workgroupSize }}u;
 const TWIDDLE_RESET: u32 = 64u;
 var<workgroup> signalTile: array<f32, N>;

 // phase drift without transcendental functions in the inner loop.
 {{ env.wgsl.resourceDeclarations }}
 // Backend-stable unit-circle twiddles for DFT kernels. Retain the phase as a
+// rational turn and route polynomial evaluation through one trigonometric core
+// for FFT, direct DFT, STFT, and cosine-window calculations.
 //
 // phase is in [0, period). sign is -1 for a forward transform and +1 for an
 // inverse transform. The result is (cos, sin).
 // Backend-stable f32 sine/cosine core.
 //
+// Shader transcendental accuracy is implementation-defined. Each path retains
+// its available phase representation, reduces it to [-pi, pi], then evaluates
+// explicit polynomials in a fixed order. The half-pi core is degree 13 for sine
+// and degree 12 for cosine.
 const PRECISE_TRIG_PI: f32 = 3.141592653589793;
 const PRECISE_TRIG_TWO_PI: f32 = 6.283185307179586;
 const PRECISE_TRIG_HALF_PI: f32 = 1.5707963267948966;
   return precise_sincos_centered(sign * PRECISE_TRIG_TWO_PI * centeredTurns);
 }
+const N: u32 = {{ transformLength }}u;
+const NUM_BINS: u32 = {{ numBins }}u;
+const WG: u32 = {{ workgroupSizeSpec }}u;
 const TWIDDLE_RESET: u32 = 64u;
 var<workgroup> signalTile: array<f32, N>;

build/webgpu/fft-radix2-dit-shared-inplace.wgsl.jinja CHANGED Viewed

@@ -2,25 +2,24 @@
 //
 // This route fills the gap between the double-buffered Stockham kernel
 // (16 bytes/element) and the global-scratch DIT fallback. A single complex
-// array needs only 8 bytes/element, extending the resident FFT range one
-// power-of-two rung beyond the 16-byte/element Stockham kernel on every device.
 // Each workgroup owns one transform; all stages remain local and use uniform
 // barriers, so only the input and requested spectrum touch global memory.
 // STFT is forward-only, so the transform sign is fixed at -1.
 {{ env.wgsl.resourceDeclarations }}
 // Backend-stable unit-circle twiddles for DFT kernels. Retain the phase as a
-// rational turn and route all polynomial evaluation through the shared precise
-// trig core so FFT, direct DFT, STFT, and cosine windows agree exactly.
 //
 // phase is in [0, period). sign is -1 for a forward transform and +1 for an
 // inverse transform. The result is (cos, sin).
 // Backend-stable f32 sine/cosine core.
 //
-// Shader transcendental accuracy is implementation-defined, and some portable
-// backends are only accurate to roughly 1e-4. Each path retains the most accurate
-// available phase representation, reduces it to [-pi, pi], then uses these
-// polynomials. The half-pi core is degree 13 for sine and degree 12 for cosine;
-// truncation error is well below one f32 ULP over its documented interval.
 const PRECISE_TRIG_PI: f32 = 3.141592653589793;
 const PRECISE_TRIG_TWO_PI: f32 = 6.283185307179586;
 const PRECISE_TRIG_HALF_PI: f32 = 1.5707963267948966;
@@ -89,12 +88,12 @@ fn dft_precise_twiddle(phase: u32, period: u32, sign: f32) -> vec2<f32> {
 }
-const N: u32 = {{ source.fftLength }}u;
 const N_HALF: u32 = N / 2u;
-const NUM_BINS: u32 = {{ source.numBins }}u;
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 const SIGN: f32 = -1.0;
-var<workgroup> fftData: array<vec2<f32>, {{ source.fftLength }}>;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wgid: vec3<u32>,
@@ -106,19 +105,19 @@ fn main(@builtin(workgroup_id) wgid: vec3<u32>,
   let start = frame * params.frameStep;
   for (var position = lid.x; position < N; position = position + WG) {
-    let n = reverseBits(position) >> (32u - {{ source.log2Length }}u);
     let signalIndex = start + n;
-    let inputBase = (batch * params.signalLength + signalIndex) * {{ source.inComp }}u;
     let w = window[n];
     var value = vec2<f32>(signal[inputBase] * w, 0.0);
-{% if source.inComp == 2 %}
     value.y = signal[inputBase + 1u] * w;
 {% endif %}
     fftData[position] = value;
   }
   workgroupBarrier();
-{% for stage in range(1, source.log2Length + 1) %}
   {
     const SPAN: u32 = 1u << {{ stage }}u;
     const HALF: u32 = SPAN >> 1u;

 //
 // This route fills the gap between the double-buffered Stockham kernel
 // (16 bytes/element) and the global-scratch DIT fallback. A single complex
+// array needs only 8 bytes/element rather than Stockham's 16, extending the
+// resident range by one power-of-two step for a fixed workgroup-storage limit.
 // Each workgroup owns one transform; all stages remain local and use uniform
 // barriers, so only the input and requested spectrum touch global memory.
 // STFT is forward-only, so the transform sign is fixed at -1.
 {{ env.wgsl.resourceDeclarations }}
 // Backend-stable unit-circle twiddles for DFT kernels. Retain the phase as a
+// rational turn and route polynomial evaluation through one trigonometric core
+// for FFT, direct DFT, STFT, and cosine-window calculations.
 //
 // phase is in [0, period). sign is -1 for a forward transform and +1 for an
 // inverse transform. The result is (cos, sin).
 // Backend-stable f32 sine/cosine core.
 //
+// Shader transcendental accuracy is implementation-defined. Each path retains
+// its available phase representation, reduces it to [-pi, pi], then evaluates
+// explicit polynomials in a fixed order. The half-pi core is degree 13 for sine
+// and degree 12 for cosine.
 const PRECISE_TRIG_PI: f32 = 3.141592653589793;
 const PRECISE_TRIG_TWO_PI: f32 = 6.283185307179586;
 const PRECISE_TRIG_HALF_PI: f32 = 1.5707963267948966;
 }
+const N: u32 = {{ fftLength }}u;
 const N_HALF: u32 = N / 2u;
+const NUM_BINS: u32 = {{ numBins }}u;
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 const SIGN: f32 = -1.0;
+var<workgroup> fftData: array<vec2<f32>, {{ fftLength }}>;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wgid: vec3<u32>,
   let start = frame * params.frameStep;
   for (var position = lid.x; position < N; position = position + WG) {
+    let n = reverseBits(position) >> (32u - {{ log2Length }}u);
     let signalIndex = start + n;
+    let inputBase = (batch * params.signalLength + signalIndex) * {{ inComp }}u;
     let w = window[n];
     var value = vec2<f32>(signal[inputBase] * w, 0.0);
+{% if inComp == 2 %}
     value.y = signal[inputBase + 1u] * w;
 {% endif %}
     fftData[position] = value;
   }
   workgroupBarrier();
+{% for stage in range(1, log2Length + 1) %}
   {
     const SPAN: u32 = 1u << {{ stage }}u;
     const HALF: u32 = SPAN >> 1u;

build/webgpu/fft-radix2-dit-storage-pass.wgsl.jinja CHANGED Viewed

@@ -1,15 +1,14 @@
-// Globally staged radix-2 DIT FFT. A single-workgroup storage FFT is efficient
-// when a device can keep a whole transform resident, but architectures that
-// serialize storage barriers or spill a large workgroup can become severely
-// under-occupied. This route exposes every transform element or butterfly to
-// the global scheduler. Separate dispatches provide the required stage barriers
-// while preserving the same bit-reversed DIT algorithm and scratch layout.
 // STFT is forward-only, so the transform sign is fixed at -1.
 {{ env.wgsl.resourceDeclarations }}
-{% if source.phase == "local_prefix" or source.phase == "stage" %}
 // Backend-stable unit-circle twiddles for DFT kernels. Retain the phase as a
-// rational turn and route all polynomial evaluation through the shared precise
-// trig core so FFT, direct DFT, STFT, and cosine windows agree exactly.
 //
 // phase is in [0, period). sign is -1 for a forward transform and +1 for an
 // inverse transform. The result is (cos, sin).
@@ -17,11 +16,10 @@
 {% set needPreciseTrigTwoPi = true %}
 // Backend-stable f32 sine/cosine core.
 //
-// Shader transcendental accuracy is implementation-defined, and some portable
-// backends are only accurate to roughly 1e-4. Each path retains the most accurate
-// available phase representation, reduces it to [-pi, pi], then uses these
-// polynomials. The half-pi core is degree 13 for sine and degree 12 for cosine;
-// truncation error is well below one f32 ULP over its documented interval.
 {% if needPreciseTrigCentered %}
 const PRECISE_TRIG_PI: f32 = 3.141592653589793;
 {% endif %}
@@ -99,62 +97,61 @@ fn dft_precise_twiddle(phase: u32, period: u32, sign: f32) -> vec2<f32> {
 {% endif %}
-const N: u32 = {{ source.fftLength }}u;
-{% if source.phase == "stage" %}
 const N_HALF: u32 = N / 2u;
 {% endif %}
-{% if source.phase == "finish" %}
-const NUM_BINS: u32 = {{ source.numBins }}u;
 {% endif %}
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
-{% if source.phase == "stage" %}
-const SPAN: u32 = 1u << {{ source.stage }}u;
 const HALF: u32 = SPAN >> 1u;
 const SIGN: f32 = -1.0;
-{% elif source.phase == "local_prefix" %}
-const LOCAL_N: u32 = {{ source.localLength }}u;
 const BLOCKS_PER_TRANSFORM: u32 = N / LOCAL_N;
 const SIGN: f32 = -1.0;
-var<workgroup> localData: array<vec2<f32>, {{ source.localLength }}>;
 {% endif %}
-{% if source.phase != "local_prefix" %}fn folded_global_index(gid: vec3<u32>, nwg: vec3<u32>) -> u32 {
-  return gid.x + gid.y * nwg.x * WG;
 }{% endif %}
 @compute @workgroup_size(WG, 1, 1)
-fn main({% if source.phase != "local_prefix" %}@builtin(global_invocation_id) gid: vec3<u32>,
         {% else %}@builtin(local_invocation_id) lid: vec3<u32>,
-        @builtin(workgroup_id) wgid: vec3<u32>,
-        {% endif %}
-        @builtin(num_workgroups) nwg: vec3<u32>) {
-{% if source.phase != "local_prefix" %}
-  let flat = folded_global_index(gid, nwg);
 {% endif %}
   let transforms = arrayLength(&scratch) / (N * 2u);
-{% if source.phase == "local_prefix" %}
-  let flatGroup = wgid.x + wgid.y * nwg.x;
   let transform = flatGroup / BLOCKS_PER_TRANSFORM;
   if (transform >= transforms) { return; }
   let block = flatGroup - transform * BLOCKS_PER_TRANSFORM;
   let position = block * LOCAL_N + lid.x;
-  let n = reverseBits(position) >> (32u - {{ source.log2Length }}u);
   let batch = transform / params.numFrames;
   let frame = transform - batch * params.numFrames;
   let start = frame * params.frameStep;
   let signalIndex = start + n;
-  let inputBase = (batch * params.signalLength + signalIndex) * {{ source.inComp }}u;
   let w = window[n];
   var value = vec2<f32>(signal[inputBase] * w, 0.0);
-{% if source.inComp == 2 %}
   value.y = signal[inputBase + 1u] * w;
 {% endif %}
   localData[lid.x] = value;
   workgroupBarrier();
-{% for stage in range(1, source.localLog2Length + 1) %}
   if (lid.x < LOCAL_N / 2u) {
     const LOCAL_SPAN_{{ stage }}: u32 = 1u << {{ stage }}u;
     const LOCAL_HALF_{{ stage }}: u32 = LOCAL_SPAN_{{ stage }} >> 1u;
@@ -179,7 +176,7 @@ fn main({% if source.phase != "local_prefix" %}@builtin(global_invocation_id) gi
   scratch[dst] = localData[lid.x].x;
   scratch[dst + 1u] = localData[lid.x].y;
-{% elif source.phase == "stage" %}
   let count = transforms * N_HALF;
   if (flat >= count) { return; }
   let transform = flat / N_HALF;

+// Globally staged radix-2 DIT FFT for transforms not admitted to the
+// workgroup-resident route. Each dispatch exposes every transform element or
+// butterfly as independent work. Dispatch boundaries provide the required
+// stage barriers while preserving the bit-reversed DIT algorithm and scratch
+// layout.
 // STFT is forward-only, so the transform sign is fixed at -1.
 {{ env.wgsl.resourceDeclarations }}
+{% if phase == "local_prefix" or phase == "stage" %}
 // Backend-stable unit-circle twiddles for DFT kernels. Retain the phase as a
+// rational turn and route polynomial evaluation through one trigonometric core
+// for FFT, direct DFT, STFT, and cosine-window calculations.
 //
 // phase is in [0, period). sign is -1 for a forward transform and +1 for an
 // inverse transform. The result is (cos, sin).
 {% set needPreciseTrigTwoPi = true %}
 // Backend-stable f32 sine/cosine core.
 //
+// Shader transcendental accuracy is implementation-defined. Each path retains
+// its available phase representation, reduces it to [-pi, pi], then evaluates
+// explicit polynomials in a fixed order. The half-pi core is degree 13 for sine
+// and degree 12 for cosine.
 {% if needPreciseTrigCentered %}
 const PRECISE_TRIG_PI: f32 = 3.141592653589793;
 {% endif %}
 {% endif %}
+const N: u32 = {{ fftLength }}u;
+{% if phase == "stage" %}
 const N_HALF: u32 = N / 2u;
 {% endif %}
+{% if phase == "finish" %}
+const NUM_BINS: u32 = {{ numBins }}u;
 {% endif %}
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+{% if phase == "stage" %}
+const SPAN: u32 = 1u << {{ stage }}u;
 const HALF: u32 = SPAN >> 1u;
 const SIGN: f32 = -1.0;
+{% elif phase == "local_prefix" %}
+const LOCAL_N: u32 = {{ localLength }}u;
 const BLOCKS_PER_TRANSFORM: u32 = N / LOCAL_N;
 const SIGN: f32 = -1.0;
+var<workgroup> localData: array<vec2<f32>, {{ localLength }}>;
 {% endif %}
+{% if phase != "local_prefix" %}fn folded_global_index(gid: vec3<u32>) -> u32 {
+  return gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
 }{% endif %}
 @compute @workgroup_size(WG, 1, 1)
+fn main({% if phase != "local_prefix" %}@builtin(global_invocation_id) gid: vec3<u32>
         {% else %}@builtin(local_invocation_id) lid: vec3<u32>,
+        @builtin(workgroup_id) wgid: vec3<u32>
+        {% endif %}) {
+{% if phase != "local_prefix" %}
+  let flat = folded_global_index(gid);
 {% endif %}
   let transforms = arrayLength(&scratch) / (N * 2u);
+{% if phase == "local_prefix" %}
+  let flatGroup = wgid.x + wgid.y * {{ DISPATCH_FOLD_WIDTH }}u;
   let transform = flatGroup / BLOCKS_PER_TRANSFORM;
   if (transform >= transforms) { return; }
   let block = flatGroup - transform * BLOCKS_PER_TRANSFORM;
   let position = block * LOCAL_N + lid.x;
+  let n = reverseBits(position) >> (32u - {{ log2Length }}u);
   let batch = transform / params.numFrames;
   let frame = transform - batch * params.numFrames;
   let start = frame * params.frameStep;
   let signalIndex = start + n;
+  let inputBase = (batch * params.signalLength + signalIndex) * {{ inComp }}u;
   let w = window[n];
   var value = vec2<f32>(signal[inputBase] * w, 0.0);
+{% if inComp == 2 %}
   value.y = signal[inputBase + 1u] * w;
 {% endif %}
   localData[lid.x] = value;
   workgroupBarrier();
+{% for stage in range(1, localLog2Length + 1) %}
   if (lid.x < LOCAL_N / 2u) {
     const LOCAL_SPAN_{{ stage }}: u32 = 1u << {{ stage }}u;
     const LOCAL_HALF_{{ stage }}: u32 = LOCAL_SPAN_{{ stage }} >> 1u;
   scratch[dst] = localData[lid.x].x;
   scratch[dst + 1u] = localData[lid.x].y;
+{% elif phase == "stage" %}
   let count = transforms * N_HALF;
   if (flat >= count) { return; }
   let transform = flat / N_HALF;

build/webgpu/fft-radix2-dit-storage.wgsl.jinja CHANGED Viewed

@@ -8,15 +8,15 @@
 // row. Scratch-backed DFT supports one-sided output. STFT accepts real or
 // complex windowed frames and copies NUM_BINS after the final stage.
 {{ env.wgsl.resourceDeclarations }}
-{% set STFT = source.mode == "stft" %}
 const PI: f32 = 3.141592653589793;
-const N: u32 = {{ source.fftLength }}u;
 const N_HALF: u32 = N / 2u;
-const LOG2_N: u32 = {{ source.log2Length }}u;
-const NUM_BINS: u32 = {{ source.numBins }}u;
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
-const SIGN: f32 = {{ "1.0" if source.inverse else "-1.0" }};
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(
@@ -35,10 +35,10 @@ fn main(
   for (var n = tid; n < N; n = n + WG) {
     var value = vec2<f32>(0.0);
     let signalIndex = start + n;
-    let inputBase = (batch * params.signalLength + signalIndex) * {{ source.inComp }}u;
     let w = window[n];
     value.x = signal[inputBase] * w;
-{% if source.inComp == 2 %}
     value.y = signal[inputBase + 1u] * w;
 {% endif %}
     let reversed = reverseBits(n) >> (32u - LOG2_N);
@@ -48,7 +48,7 @@ fn main(
   }
   storageBarrier();
-{% for stage in range(1, source.log2Length + 1) %}
   {
     const SPAN: u32 = 1u << {{ stage }}u;
     const HALF: u32 = SPAN >> 1u;

 // row. Scratch-backed DFT supports one-sided output. STFT accepts real or
 // complex windowed frames and copies NUM_BINS after the final stage.
 {{ env.wgsl.resourceDeclarations }}
+{% set STFT = modeSpec == "stft" %}
 const PI: f32 = 3.141592653589793;
+const N: u32 = {{ fftLength }}u;
 const N_HALF: u32 = N / 2u;
+const LOG2_N: u32 = {{ log2Length }}u;
+const NUM_BINS: u32 = {{ numBins }}u;
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+const SIGN: f32 = {{ "1.0" if inverse else "-1.0" }};
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(
   for (var n = tid; n < N; n = n + WG) {
     var value = vec2<f32>(0.0);
     let signalIndex = start + n;
+    let inputBase = (batch * params.signalLength + signalIndex) * {{ inComp }}u;
     let w = window[n];
     value.x = signal[inputBase] * w;
+{% if inComp == 2 %}
     value.y = signal[inputBase + 1u] * w;
 {% endif %}
     let reversed = reverseBits(n) >> (32u - LOG2_N);
   }
   storageBarrier();
+{% for stage in range(1, log2Length + 1) %}
   {
     const SPAN: u32 = 1u << {{ stage }}u;
     const HALF: u32 = SPAN >> 1u;

build/webgpu/fft-real-packed-stockham-shared.wgsl.jinja CHANGED Viewed

@@ -4,12 +4,12 @@
 // complex transform, then reconstruct the N/2+1 non-redundant bins. This cuts
 // the Stockham butterfly count almost in half while preserving the same fused
 // window load as the full complex STFT kernel.
-const N: u32 = {{ source.fftLength }}u;
 const PACKED_N: u32 = N / 2u;
 const PACKED_HALF: u32 = PACKED_N / 2u;
 const NUM_BINS: u32 = PACKED_N + 1u;
 const FFT_LANES: u32 = {{ tunables.WORKGROUP_SIZE }}u;
-const TRANSFORMS_PER_WG: u32 = {{ source.transformsPerWorkgroup | default(1) }}u;
 const WG: u32 = FFT_LANES * TRANSFORMS_PER_WG;
 const STAGE_TWIDDLES: u32 = PACKED_N - 1u;
 const SHARED_STRIDE: u32 = 2u * PACKED_N;

 // complex transform, then reconstruct the N/2+1 non-redundant bins. This cuts
 // the Stockham butterfly count almost in half while preserving the same fused
 // window load as the full complex STFT kernel.
+const N: u32 = {{ fftLength }}u;
 const PACKED_N: u32 = N / 2u;
 const PACKED_HALF: u32 = PACKED_N / 2u;
 const NUM_BINS: u32 = PACKED_N + 1u;
 const FFT_LANES: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+const TRANSFORMS_PER_WG: u32 = {{ transformsPerWorkgroup | default(1) }}u;
 const WG: u32 = FFT_LANES * TRANSFORMS_PER_WG;
 const STAGE_TWIDDLES: u32 = PACKED_N - 1u;
 const SHARED_STRIDE: u32 = 2u * PACKED_N;

build/webgpu/fft-stockham-shared.wgsl.jinja CHANGED Viewed

@@ -1,11 +1,12 @@
 {{ env.wgsl.resourceDeclarations }}
-{% set MIXED_TAIL = source.mixedRadixTail if source.mixedRadixTail is defined else false %}
-{% set USE_TWIDDLE_TABLE = source.useTwiddleTable if source.useTwiddleTable is defined else false %}
 {% set NEED_PRECISE_TWIDDLE = not USE_TWIDDLE_TABLE or MIXED_TAIL %}
 {% if NEED_PRECISE_TWIDDLE %}
 // Backend-stable unit-circle twiddles for DFT kernels. Retain the phase as a
-// rational turn and route all polynomial evaluation through the shared precise
-// trig core so FFT, direct DFT, STFT, and cosine windows agree exactly.
 //
 // phase is in [0, period). sign is -1 for a forward transform and +1 for an
 // inverse transform. The result is (cos, sin).
@@ -13,11 +14,10 @@
 {% set needPreciseTrigTwoPi = true %}
 // Backend-stable f32 sine/cosine core.
 //
-// Shader transcendental accuracy is implementation-defined, and some portable
-// backends are only accurate to roughly 1e-4. Each path retains the most accurate
-// available phase representation, reduces it to [-pi, pi], then uses these
-// polynomials. The half-pi core is degree 13 for sine and degree 12 for cosine;
-// truncation error is well below one f32 ULP over its documented interval.
 {% if needPreciseTrigCentered %}
 const PRECISE_TRIG_PI: f32 = 3.141592653589793;
 {% endif %}
@@ -96,49 +96,42 @@ fn dft_precise_twiddle(phase: u32, period: u32, sign: f32) -> vec2<f32> {
 {% endif %}
 {% set TAIL_ANCHOR = 32 %}
 {% set REANCHOR_TAIL = MIXED_TAIL
-    and source.fftLength / source.radix2Length > TAIL_ANCHOR %}
 // Radix-2 Stockham autosort FFT: one transform per workgroup, entirely in
 // shared memory. log2(N) ping-pong stages alternate between the two halves of
 // a 2*N scratch pair (no bit reversal; the autosort write pattern leaves the
 // final stage in natural order). Threads stride over the N/2 butterflies of
-// each stage with a workgroupBarrier between stages. Large transform batches
-// may bind a compact precomputed twiddle table; smaller batches evaluate the
-// same roots inline and avoid the extra dispatch.
-// Mode "dft" views the input as [outer, N, inner, inComp] and the output as
-// [outer, numBins, inner, 2]; params.inner carries the (runtime) inner count
-// and dispatch.x = outer * inner workgroups. Mode "stft" runs one workgroup
-// per (batch, frame): the window multiply is fused into the frame load. STFT
-// selection admits only complete frames, so every frame load is in bounds. The
-// inverse+onesided (IRFFT) combination is outside this specialization.
-// Mode "logmel" runs one workgroup per spectrogram frame of
-// a real mono signal: the frame load is centered (start - N/2) with reflection
-// below zero and zeros at/after params.audioLength, and the epilogue folds the
-// one-sided power spectrum through packed triangular mel bands into log10
-// mel-major output instead of storing complex bins.
-const N: u32 = {{ source.fftLength }}u;
 const N_HALF: u32 = N / 2u;
 {% if MIXED_TAIL %}
-const RADIX2_N: u32 = {{ source.radix2Length }}u;
 const RADIX2_HALF: u32 = RADIX2_N / 2u;
 const TAIL_RADIX: u32 = N / RADIX2_N;
 {% if REANCHOR_TAIL %}
 const TAIL_ANCHOR: u32 = {{ TAIL_ANCHOR }}u;
 {% endif %}
 {% endif %}
-const NUM_BINS: u32 = {{ source.numBins }}u;
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 {% if NEED_PRECISE_TWIDDLE %}
 // Twiddle = exp(SIGN * i * 2*pi * k / span); SIGN flips for the inverse
-// transform, exactly like the naive kernels' `sign` select.
-const SIGN: f32 = {{ "1.0" if source.inverse else "-1.0" }};
 {% endif %}
 {% if MIXED_TAIL %}
-// Walk consecutive powers of one unit root in each mixed-radix tail. A single
-// backend-stable root evaluation per output bin replaces one precise
-// polynomial evaluation per tail residue for every supported factorization.
 fn complex_mul(a: vec2<f32>, b: vec2<f32>) -> vec2<f32> {
   return vec2<f32>(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
 }
@@ -153,11 +146,8 @@ fn main(@builtin(workgroup_id) wgid: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>,
         @builtin(num_workgroups) nwg: vec3<u32>) {
   let tid = lid.x;
-  // Flat workgroup index, 2D-folded: dispatches that would exceed the
-  // maxComputeWorkgroupsPerDimension limit on a single axis split the
-  // one-workgroup-per-(batch,frame) [or per-(outer,inner)] grid across x/y and
-  // recover the flat index here. Collapses to wgid.x when nwg.y == 1 (every DFT
-  // dispatch and any STFT dispatch under the per-axis limit).
   let flat = wgid.x + wgid.y * nwg.x;
   let batch = flat / params.numFrames;
@@ -172,14 +162,14 @@ fn main(@builtin(workgroup_id) wgid: vec3<u32>,
     let sample_n = n;
 {% endif %}
     let signal_index = start + sample_n;
-    let base = (batch * params.signalLength + signal_index) * {{ source.inComp }}u;
-{% if source.hasWindow is defined and not source.hasWindow %}
     let w = 1.0;
 {% else %}
     let w = window[sample_n];
 {% endif %}
     let xr = signal[base] * w;
-{% if source.inComp == 2 %}
     let xi = signal[base + 1u] * w;
 {% else %}
     let xi = 0.0;
@@ -193,8 +183,7 @@ fn main(@builtin(workgroup_id) wgid: vec3<u32>,
   // applies w = exp(SIGN*i*pi*(j mod ns)/ns) to b, and writes
   //   dst[(j div ns)*2*ns + (j mod ns)]      = a + w*b
   //   dst[(j div ns)*2*ns + (j mod ns) + ns] = a - w*b
-  // |angle| <= pi, so twiddle accuracy never depends on large-argument
-  // sin/cos reduction (unlike the naive kernels' k*n angles).
   var srcOff = 0u;
   var dstOff = N;
 {% if MIXED_TAIL %}
@@ -254,8 +243,7 @@ fn main(@builtin(workgroup_id) wgid: vec3<u32>,
   }
   // The final swap leaves the natural-order result at srcOff. The 1/N inverse
-  // scale is a power-of-two reciprocal, so the multiply is exact and matches
-  // the naive kernels' division.
   let outRow = (batch * params.numFrames + frame) * NUM_BINS;
   for (var k = tid; k < NUM_BINS; k = k + WG) {
     let outBase = (outRow + k) * 2u;

 {{ env.wgsl.resourceDeclarations }}
+{% set MIXED_TAIL = mixedRadixTail if mixedRadixTail is defined else false %}
+{% set USE_TWIDDLE_TABLE = useTwiddleTable if useTwiddleTable is defined else false %}
+{% set radix2Length = radix2Length | default(fftLength) %}
 {% set NEED_PRECISE_TWIDDLE = not USE_TWIDDLE_TABLE or MIXED_TAIL %}
 {% if NEED_PRECISE_TWIDDLE %}
 // Backend-stable unit-circle twiddles for DFT kernels. Retain the phase as a
+// rational turn and route polynomial evaluation through one trigonometric core
+// for FFT, direct DFT, STFT, and cosine-window calculations.
 //
 // phase is in [0, period). sign is -1 for a forward transform and +1 for an
 // inverse transform. The result is (cos, sin).
 {% set needPreciseTrigTwoPi = true %}
 // Backend-stable f32 sine/cosine core.
 //
+// Shader transcendental accuracy is implementation-defined. Each path retains
+// its available phase representation, reduces it to [-pi, pi], then evaluates
+// explicit polynomials in a fixed order. The half-pi core is degree 13 for sine
+// and degree 12 for cosine.
 {% if needPreciseTrigCentered %}
 const PRECISE_TRIG_PI: f32 = 3.141592653589793;
 {% endif %}
 {% endif %}
 {% set TAIL_ANCHOR = 32 %}
 {% set REANCHOR_TAIL = MIXED_TAIL
+    and fftLength / radix2Length > TAIL_ANCHOR %}
 // Radix-2 Stockham autosort FFT: one transform per workgroup, entirely in
 // shared memory. log2(N) ping-pong stages alternate between the two halves of
 // a 2*N scratch pair (no bit reversal; the autosort write pattern leaves the
 // final stage in natural order). Threads stride over the N/2 butterflies of
+// each stage with a workgroupBarrier between stages.
+{% if USE_TWIDDLE_TABLE %}
+// This specialization reads roots from a compact precomputed twiddle table.
+{% else %}
+// This specialization evaluates roots inline.
+{% endif %}
+// One workgroup owns each (batch, frame). The window multiplication is fused
+// into an in-bounds frame load, and the output stores numBins complex values.
+const N: u32 = {{ fftLength }}u;
 const N_HALF: u32 = N / 2u;
 {% if MIXED_TAIL %}
+const RADIX2_N: u32 = {{ radix2Length }}u;
 const RADIX2_HALF: u32 = RADIX2_N / 2u;
 const TAIL_RADIX: u32 = N / RADIX2_N;
 {% if REANCHOR_TAIL %}
 const TAIL_ANCHOR: u32 = {{ TAIL_ANCHOR }}u;
 {% endif %}
 {% endif %}
+const NUM_BINS: u32 = {{ numBins }}u;
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 {% if NEED_PRECISE_TWIDDLE %}
 // Twiddle = exp(SIGN * i * 2*pi * k / span); SIGN flips for the inverse
+// transform.
+const SIGN: f32 = {{ "1.0" if inverse else "-1.0" }};
 {% endif %}
 {% if MIXED_TAIL %}
+// Walk consecutive powers of one unit root in each mixed-radix tail. One
+// precise root evaluation per output bin anchors the recurrence.
 fn complex_mul(a: vec2<f32>, b: vec2<f32>) -> vec2<f32> {
   return vec2<f32>(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
 }
         @builtin(local_invocation_id) lid: vec3<u32>,
         @builtin(num_workgroups) nwg: vec3<u32>) {
   let tid = lid.x;
+  // Flat workgroup index for a grid folded across x and y. This reduces to
+  // wgid.x when the dispatch uses one row.
   let flat = wgid.x + wgid.y * nwg.x;
   let batch = flat / params.numFrames;
     let sample_n = n;
 {% endif %}
     let signal_index = start + sample_n;
+    let base = (batch * params.signalLength + signal_index) * {{ inComp }}u;
+{% if hasWindow is defined and not hasWindow %}
     let w = 1.0;
 {% else %}
     let w = window[sample_n];
 {% endif %}
     let xr = signal[base] * w;
+{% if inComp == 2 %}
     let xi = signal[base + 1u] * w;
 {% else %}
     let xi = 0.0;
   // applies w = exp(SIGN*i*pi*(j mod ns)/ns) to b, and writes
   //   dst[(j div ns)*2*ns + (j mod ns)]      = a + w*b
   //   dst[(j div ns)*2*ns + (j mod ns) + ns] = a - w*b
+  // |angle| <= pi, which bounds the argument supplied to sin and cos.
   var srcOff = 0u;
   var dstOff = N;
 {% if MIXED_TAIL %}
   }
   // The final swap leaves the natural-order result at srcOff. The 1/N inverse
+  // scale is an exactly representable power-of-two reciprocal applied once.
   let outRow = (batch * params.numFrames + frame) * NUM_BINS;
   for (var k = tid; k < NUM_BINS; k = k + WG) {
     let outBase = (outRow + k) * 2u;

build/webgpu/fft-stockham-twiddles.wgsl.jinja CHANGED Viewed

@@ -1,17 +1,16 @@
 {{ env.wgsl.resourceDeclarations }}
 // Backend-stable unit-circle twiddles for DFT kernels. Retain the phase as a
-// rational turn and route all polynomial evaluation through the shared precise
-// trig core so FFT, direct DFT, STFT, and cosine windows agree exactly.
 //
 // phase is in [0, period). sign is -1 for a forward transform and +1 for an
 // inverse transform. The result is (cos, sin).
 // Backend-stable f32 sine/cosine core.
 //
-// Shader transcendental accuracy is implementation-defined, and some portable
-// backends are only accurate to roughly 1e-4. Each path retains the most accurate
-// available phase representation, reduces it to [-pi, pi], then uses these
-// polynomials. The half-pi core is degree 13 for sine and degree 12 for cosine;
-// truncation error is well below one f32 ULP over its documented interval.
 const PRECISE_TRIG_PI: f32 = 3.141592653589793;
 const PRECISE_TRIG_TWO_PI: f32 = 6.283185307179586;
 const PRECISE_TRIG_HALF_PI: f32 = 1.5707963267948966;
@@ -86,9 +85,9 @@ fn dft_precise_twiddle(phase: u32, period: u32, sign: f32) -> vec2<f32> {
 // each root once avoids repeating sin/cos for every frame while preserving the
 // same f32 twiddle values and butterfly arithmetic.
 // STFT is forward-only, so the transform sign is fixed at -1.
-const N: u32 = {{ source.fftLength }}u;
 const SIGN: f32 = -1.0;
-{% if source.realPacked %}
 const PACKED_N: u32 = N / 2u;
 const STAGE_TWIDDLES: u32 = PACKED_N - 1u;
@@ -96,7 +95,7 @@ const STAGE_TWIDDLES: u32 = PACKED_N - 1u;
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   let tableIndex = gid.x;
-{% if source.realPacked %}
   if (tableIndex >= N) {
     return;
   }

 {{ env.wgsl.resourceDeclarations }}
 // Backend-stable unit-circle twiddles for DFT kernels. Retain the phase as a
+// rational turn and route polynomial evaluation through one trigonometric core
+// for FFT, direct DFT, STFT, and cosine-window calculations.
 //
 // phase is in [0, period). sign is -1 for a forward transform and +1 for an
 // inverse transform. The result is (cos, sin).
 // Backend-stable f32 sine/cosine core.
 //
+// Shader transcendental accuracy is implementation-defined. Each path retains
+// its available phase representation, reduces it to [-pi, pi], then evaluates
+// explicit polynomials in a fixed order. The half-pi core is degree 13 for sine
+// and degree 12 for cosine.
 const PRECISE_TRIG_PI: f32 = 3.141592653589793;
 const PRECISE_TRIG_TWO_PI: f32 = 6.283185307179586;
 const PRECISE_TRIG_HALF_PI: f32 = 1.5707963267948966;
 // each root once avoids repeating sin/cos for every frame while preserving the
 // same f32 twiddle values and butterfly arithmetic.
 // STFT is forward-only, so the transform sign is fixed at -1.
+const N: u32 = {{ fftLength }}u;
 const SIGN: f32 = -1.0;
+{% if realPacked %}
 const PACKED_N: u32 = N / 2u;
 const STAGE_TWIDDLES: u32 = PACKED_N - 1u;
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   let tableIndex = gid.x;
+{% if realPacked %}
   if (tableIndex >= N) {
     return;
   }

build/webgpu/manifest.json CHANGED Viewed

@@ -2,61 +2,27 @@
   "domain": "ai.onnx",
   "name": "STFT",
   "sinceVersion": 17,
-  "description": "Computes the short-time Fourier transform (STFT) of a rank-3 real or complex signal by sliding a complete frame over the input and applying the DFT. The `frameStep` scalar represents the required positive `frame_step` input. At least one of `window` or `frameLength` is required; when both are supplied, their lengths must match. The output contains the real and imaginary components of every frequency bin for each complete frame.",
-  "inputs": [
-    {
-      "role": "signal",
-      "dtype": "T",
-      "rank": 3,
-      "description": "Real input signal with shape `(batch_size, signal_length, 1)`, or complex input with shape `(batch_size, signal_length, 2)`."
-    },
-    {
-      "role": "window",
-      "dtype": "T",
-      "rank": 1,
-      "optional": true,
-      "description": "Optional 1-D window tensor with shape `(window_length)` to multiply by each frame before the DFT."
-    }
-  ],
-  "outputs": [
-    {
-      "role": "output",
       "dtype": "T",
       "rank": 4,
-      "description": "STFT result with shape `(batch_size, frames, dft_unique_bins, 2)` in one-sided mode or `(batch_size, frames, frame_length, 2)` in two-sided mode, with real and imaginary parts in the last dimension. `frames = floor((signal_length - frame_length) / frame_step) + 1`; partial trailing frames are not emitted.",
       "shape": ["dim(shapes.signal, 0)", "expectedFrames", "floor(frameLength / 2) + 1 if attrs.onesided == 1 else frameLength", 2]
     }
-  ],
-  "attributes": { "onesided": 1 },
-  "attributeDescriptions": {
-    "onesided": "When 1 (default), returns `floor(n_fft / 2) + 1` non-redundant frequency bins with indices `[0, ..., floor(n_fft / 2)]` for a real-valued input (RFFT); set to 0 to return the full spectrum."
   },
-  "attributeConstraints": { "onesided": { "values": [0, 1] } },
-  "typeConstraints": { "T": ["float32"] },
   "args": {
-    "signal": { "kind": "tensor", "semantic": "signal", "role": "input" },
-    "window": { "kind": "tensor", "semantic": "window", "role": "input", "required": false },
-    "output": { "kind": "tensor", "semantic": "output", "role": "output" },
-    "frameStep": {
-      "kind": "i32",
-      "semantic": "frame_step",
-      "role": "attribute",
-      "description": "Positive number of signal positions between the starts of consecutive frames."
-    },
-    "frameLength": {
-      "kind": "i32",
-      "semantic": "frame_length",
-      "role": "attribute",
-      "required": false,
-      "description": "Optional number of signal values in each frame; must equal the window length when both are supplied."
-    }
   },
   "tunables": {
-    "WORKGROUP_SIZE": 128,
-    "TWIDDLE_TABLE_MIN_TRANSFORMS": 4096,
-    "TILED_REAL_MIN_LENGTH": 512,
-    "GLOBAL_STAGE_LOCAL_LENGTH": 256,
-    "GLOBAL_STAGE_MAX_DEVICE_STORAGE_BYTES": 16384
   },
   "derive": {
     "windowLength": "dim(shapes.window, 0) if (present.window and ranks.window == 1) else 0",
@@ -67,7 +33,7 @@
     "rank3OutputLayout": "ranks.signal == 3 and ranks.output == 4 and dim(shapes.output, 0) == dim(shapes.signal, 0) and dim(shapes.output, 1) == expectedFrames and dim(shapes.signal, 2) >= 1 and dim(shapes.signal, 2) <= 2 and dim(shapes.output, 3) == 2",
     "rank3WindowLayout": "present.window and ranks.window == 1 and rank3OutputLayout and frameControlsValid",
     "rank3WindowSpectrum": "false if not present.window else ((attrs.onesided == 0 and dim(shapes.output, 2) == dim(shapes.window, 0)) or (attrs.onesided != 0 and dim(shapes.signal, 2) == 1 and dim(shapes.output, 2) == floor(dim(shapes.window, 0) / 2) + 1))",
-    "rank3WindowDispatchFits": "dim(shapes.output, 1) <= device.limits.maxComputeWorkgroupsPerDimension and dim(shapes.output, 0) <= device.limits.maxComputeWorkgroupsPerDimension",
     "rank3WindowSharedFftFits": "false if not present.window else (rank3WindowDispatchFits and pow2ceil(dim(shapes.window, 0)) == dim(shapes.window, 0) and dim(shapes.window, 0) >= 4 and dim(shapes.window, 0) * 16 <= device.limits.maxComputeWorkgroupStorageSize)",
     "rank3WindowFftRadix2": "1 if not present.window else pick([[dim(shapes.window, 0) % 256 == 0, 256], [dim(shapes.window, 0) % 128 == 0, 128], [dim(shapes.window, 0) % 64 == 0, 64], [dim(shapes.window, 0) % 32 == 0, 32], [dim(shapes.window, 0) % 16 == 0, 16], [dim(shapes.window, 0) % 8 == 0, 8], [dim(shapes.window, 0) % 4 == 0, 4]], 1)",
     "rank3WindowFftTailRadix": "1 if not present.window else dim(shapes.window, 0) / rank3WindowFftRadix2",
@@ -76,235 +42,54 @@
     "rank3WindowRadix2ScratchFits": "false if not present.window else (rank3WindowDispatchFits and pow2ceil(dim(shapes.window, 0)) == dim(shapes.window, 0) and dim(shapes.window, 0) >= 4 and dim(shapes.window, 0) * 16 > device.limits.maxComputeWorkgroupStorageSize and dim(shapes.output, 0) * dim(shapes.output, 1) > 0 and dim(shapes.output, 0) * dim(shapes.output, 1) * dim(shapes.window, 0) * 2 * 4 <= device.limits.maxStorageBufferBindingSize and dim(shapes.output, 0) * dim(shapes.output, 1) * dim(shapes.window, 0) * 2 * 4 <= device.limits.maxBufferSize)",
     "rank3WindowInplaceFftFits": "false if not present.window else (rank3WindowDispatchFits and pow2ceil(dim(shapes.window, 0)) == dim(shapes.window, 0) and dim(shapes.window, 0) >= 4 and dim(shapes.window, 0) * 16 > device.limits.maxComputeWorkgroupStorageSize and dim(shapes.window, 0) * 8 <= device.limits.maxComputeWorkgroupStorageSize and dim(shapes.output, 0) * dim(shapes.output, 1) > 0)",
     "rank3WindowGlobalFftFits": "false if not present.window else (rank3WindowRadix2ScratchFits and dim(shapes.window, 0) >= 2 * tunables.GLOBAL_STAGE_LOCAL_LENGTH and dim(shapes.window, 0) * 8 > device.limits.maxComputeWorkgroupStorageSize)",
-    "rank3WindowTiledRealFits": "false if not present.window else (rank3WindowDispatchFits and ranks.signal == 3 and dim(shapes.signal, 2) == 1 and dim(shapes.window, 0) >= tunables.TILED_REAL_MIN_LENGTH and dim(shapes.window, 0) * 4 <= device.limits.maxComputeWorkgroupStorageSize and dim(shapes.output, 0) > 0 and dim(shapes.output, 1) > 0 and ceilDiv(dim(shapes.output, 2), 256) <= device.limits.maxComputeWorkgroupsPerDimension)",
     "packedFourWayFits": "false if not present.window else (device.limits.maxComputeInvocationsPerWorkgroup >= 256 and device.limits.maxComputeWorkgroupSizeX >= 256 and dim(shapes.window, 0) * 32 <= device.limits.maxComputeWorkgroupStorageSize)",
     "packedTransformsPerWorkgroup": "4 if packedFourWayFits else 1",
     "rank3NoWindowLayout": "not present.window and args.frameLength is defined and rank3OutputLayout and frameControlsValid",
     "rank3NoWindowSpectrum": "(attrs.onesided == 0 and dim(shapes.output, 2) == frameLength) or (attrs.onesided != 0 and dim(shapes.signal, 2) == 1 and dim(shapes.output, 2) == floor(frameLength / 2) + 1)",
-    "rank3NoWindowSharedFftFits": "false if ranks.signal != 3 else ((attrs.onesided == 0 or (dim(shapes.signal, 2) == 1 and dim(shapes.output, 2) >= 3)) and pow2ceil(frameLength) == frameLength and frameLength >= 4 and frameLength * 16 <= device.limits.maxComputeWorkgroupStorageSize and dim(shapes.output, 1) <= device.limits.maxComputeWorkgroupsPerDimension and dim(shapes.output, 0) <= device.limits.maxComputeWorkgroupsPerDimension)"
   },
-  "bindingSets": {
-    "stftWindowFft": [
-      {
-        "name": "signal",
-        "arg": "signal",
-        "semantic": "signal",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "window",
-        "arg": "window",
-        "semantic": "window",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "output", "arg": "output", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "signalLength", "type": "u32", "value": "dim(shapes.signal, 1)" },
-            { "name": "numFrames", "type": "u32", "value": "dim(shapes.output, 1)" },
-            { "name": "frameStep", "type": "u32", "value": "args.frameStep" }
-          ]
-        }
-      }
-    ],
-    "stftTwiddleWrite": [
-      { "name": "twiddles", "semantic": "fft_twiddles", "buffer": { "type": "storage" }, "elementType": "f32" }
-    ],
-    "stftWindowFftTwiddles": [
-      {
-        "name": "signal",
-        "arg": "signal",
-        "semantic": "signal",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "window",
-        "arg": "window",
-        "semantic": "window",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "twiddles",
-        "semantic": "fft_twiddles",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "output", "arg": "output", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "signalLength", "type": "u32", "value": "dim(shapes.signal, 1)" },
-            { "name": "numFrames", "type": "u32", "value": "dim(shapes.output, 1)" },
-            { "name": "frameStep", "type": "u32", "value": "args.frameStep" }
-          ]
-        }
-      }
-    ],
-    "stftNoWindowFft": [
-      {
-        "name": "signal",
-        "arg": "signal",
-        "semantic": "signal",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "output", "arg": "output", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "signalLength", "type": "u32", "value": "dim(shapes.signal, 1)" },
-            { "name": "numFrames", "type": "u32", "value": "dim(shapes.output, 1)" },
-            { "name": "frameStep", "type": "u32", "value": "args.frameStep" }
-          ]
-        }
-      }
-    ],
-    "stftWindowScratch": [
-      {
-        "name": "signal",
-        "arg": "signal",
-        "semantic": "signal",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "window",
-        "arg": "window",
-        "semantic": "window",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "scratch", "semantic": "fft_scratch", "buffer": { "type": "storage" }, "elementType": "f32" },
-      { "name": "output", "arg": "output", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "signalLength", "type": "u32", "value": "dim(shapes.signal, 1)" },
-            { "name": "numFrames", "type": "u32", "value": "dim(shapes.output, 1)" },
-            { "name": "frameStep", "type": "u32", "value": "args.frameStep" }
-          ]
-        }
-      }
-    ],
-    "stftWindowScratchLocalPrefix": [
-      {
-        "name": "signal",
-        "arg": "signal",
-        "semantic": "signal",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "window",
-        "arg": "window",
-        "semantic": "window",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "scratch", "semantic": "fft_scratch", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "signalLength", "type": "u32", "value": "dim(shapes.signal, 1)" },
-            { "name": "numFrames", "type": "u32", "value": "dim(shapes.output, 1)" },
-            { "name": "frameStep", "type": "u32", "value": "args.frameStep" }
-          ]
-        }
-      }
-    ],
-    "stftWindowScratchStage": [
-      { "name": "scratch", "semantic": "fft_scratch", "buffer": { "type": "storage" }, "elementType": "f32" }
-    ],
-    "stftWindowScratchFinish": [
-      { "name": "scratch", "semantic": "fft_scratch", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
-      { "name": "output", "arg": "output", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "f32" }
-    ],
-    "stftWindowGeneral": [
-      {
-        "name": "signal",
-        "arg": "signal",
-        "semantic": "signal",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "window",
-        "arg": "window",
-        "semantic": "window",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "output", "arg": "output", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "count", "type": "u32", "value": "numel(shapes.output) / 2" },
-            { "name": "signalLength", "type": "u32", "value": "dim(shapes.signal, 1)" },
-            { "name": "signalComponents", "type": "u32", "value": "dim(shapes.signal, 2)" },
-            { "name": "frameLength", "type": "u32", "value": "dim(shapes.window, 0)" },
-            { "name": "frames", "type": "u32", "value": "dim(shapes.output, 1)" },
-            { "name": "bins", "type": "u32", "value": "dim(shapes.output, 2)" },
-            { "name": "frameStep", "type": "u32", "value": "args.frameStep" }
-          ]
-        }
-      }
-    ],
-    "stftNoWindowGeneral": [
-      {
-        "name": "signal",
-        "arg": "signal",
-        "semantic": "signal",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "output", "arg": "output", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "count", "type": "u32", "value": "numel(shapes.output) / 2" },
-            { "name": "signalLength", "type": "u32", "value": "dim(shapes.signal, 1)" },
-            { "name": "signalComponents", "type": "u32", "value": "dim(shapes.signal, 2)" },
-            { "name": "frameLength", "type": "u32", "value": "args.frameLength" },
-            { "name": "frames", "type": "u32", "value": "dim(shapes.output, 1)" },
-            { "name": "bins", "type": "u32", "value": "dim(shapes.output, 2)" },
-            { "name": "frameStep", "type": "u32", "value": "args.frameStep" }
-          ]
-        }
-      }
-    ]
   },
   "variants": [
     {
@@ -315,18 +100,16 @@
         {
           "id": "main",
           "name": "STFT.NoWindowFftShared",
-          "source": {
-            "shader": "fft-stockham-shared.wgsl.jinja",
-            "inputs": {
-              "mode": "\"stft\"",
-              "fftLength": "frameLength",
-              "numBins": "dim(shapes.output, 2)",
-              "inComp": "dim(shapes.signal, 2)",
-              "inverse": false,
-              "hasWindow": false
-            }
           },
-          "bindings": "stftNoWindowFft",
           "dispatch": { "x": "dim(shapes.output, 1)", "y": "dim(shapes.output, 0)" }
         }
       ]
@@ -334,23 +117,21 @@
     {
       "id": "rank3_window_fft_shared_inplace",
       "priority": 32,
-      "requires": { "limits": { "maxComputeInvocationsPerWorkgroup": 256, "maxComputeWorkgroupSizeX": 256 } },
-      "tunables": { "WORKGROUP_SIZE": 256 },
       "when": ["rank3WindowLayout", "rank3WindowSpectrum", "rank3WindowInplaceFftFits"],
       "passes": [
         {
           "id": "main",
           "name": "STFT.FftSharedInplace",
-          "source": {
-            "shader": "fft-radix2-dit-shared-inplace.wgsl.jinja",
-            "inputs": {
-              "fftLength": "dim(shapes.window, 0)",
-              "log2Length": "rank3WindowFftLog2",
-              "numBins": "dim(shapes.output, 2)",
-              "inComp": "dim(shapes.signal, 2)"
-            }
           },
-          "bindings": "stftWindowFft",
           "dispatch": { "x": "dim(shapes.output, 1)", "y": "dim(shapes.output, 0)" }
         }
       ]
@@ -358,10 +139,10 @@
     {
       "id": "rank3_window_fft_global_stages",
       "priority": 31,
-      "requires": { "limits": { "maxComputeInvocationsPerWorkgroup": 256, "maxComputeWorkgroupSizeX": 256 } },
-      "tunables": { "WORKGROUP_SIZE": 256 },
-      "demoteWhen": ["device.limits.maxComputeWorkgroupStorageSize > tunables.GLOBAL_STAGE_MAX_DEVICE_STORAGE_BYTES"],
       "when": ["rank3WindowLayout", "rank3WindowSpectrum", "rank3WindowGlobalFftFits"],
       "intermediates": [
         {
           "id": "fft_scratch",
@@ -373,21 +154,21 @@
         {
           "id": "local_prefix",
           "name": "STFT.FftLocalPrefix",
-          "source": {
-            "shader": "fft-radix2-dit-storage-pass.wgsl.jinja",
-            "inputs": {
-              "phase": "\"local_prefix\"",
-              "fftLength": "dim(shapes.window, 0)",
-              "log2Length": "rank3WindowFftLog2",
-              "localLength": "tunables.GLOBAL_STAGE_LOCAL_LENGTH",
-              "localLog2Length": "log2ceil(tunables.GLOBAL_STAGE_LOCAL_LENGTH)",
-              "numBins": "dim(shapes.output, 2)",
-              "inComp": "dim(shapes.signal, 2)"
-            }
           },
-          "bindings": "stftWindowScratchLocalPrefix",
           "dispatch": {
-            "workgroups": "dim(shapes.output, 0) * dim(shapes.output, 1) * dim(shapes.window, 0) / tunables.GLOBAL_STAGE_LOCAL_LENGTH"
           }
         },
         {
@@ -397,21 +178,20 @@
             {
               "id": "butterflies",
               "name": "STFT.FftGlobalStage",
-              "source": {
-                "shader": "fft-radix2-dit-storage-pass.wgsl.jinja",
-                "inputs": {
-                  "phase": "\"stage\"",
-                  "stage": "repeat.stage + log2ceil(tunables.GLOBAL_STAGE_LOCAL_LENGTH) + 1",
-                  "fftLength": "dim(shapes.window, 0)",
-                  "log2Length": "rank3WindowFftLog2",
-                  "numBins": "dim(shapes.output, 2)",
-                  "inComp": "dim(shapes.signal, 2)"
-                }
               },
-              "bindings": "stftWindowScratchStage",
               "dispatch": {
-                "threads": "dim(shapes.output, 0) * dim(shapes.output, 1) * dim(shapes.window, 0) / 2",
-                "workgroupSize": "tunables.WORKGROUP_SIZE"
               }
             }
           ]
@@ -419,20 +199,22 @@
         {
           "id": "finish",
           "name": "STFT.FftGlobalFinish",
-          "source": {
-            "shader": "fft-radix2-dit-storage-pass.wgsl.jinja",
-            "inputs": {
-              "phase": "\"finish\"",
-              "fftLength": "dim(shapes.window, 0)",
-              "log2Length": "rank3WindowFftLog2",
-              "numBins": "dim(shapes.output, 2)",
-              "inComp": "dim(shapes.signal, 2)"
-            }
           },
-          "bindings": "stftWindowScratchFinish",
           "dispatch": {
-            "threads": "dim(shapes.output, 0) * dim(shapes.output, 1) * dim(shapes.output, 2)",
-            "workgroupSize": "tunables.WORKGROUP_SIZE"
           }
         }
       ]
@@ -452,18 +234,16 @@
         {
           "id": "main",
           "name": "STFT.FftStorageInplace",
-          "source": {
-            "shader": "fft-radix2-dit-storage.wgsl.jinja",
-            "inputs": {
-              "mode": "\"stft\"",
-              "fftLength": "dim(shapes.window, 0)",
-              "log2Length": "rank3WindowFftLog2",
-              "numBins": "dim(shapes.output, 2)",
-              "inComp": "dim(shapes.signal, 2)",
-              "inverse": false
-            }
           },
-          "bindings": "stftWindowScratch",
           "dispatch": { "x": "dim(shapes.output, 1)", "y": "dim(shapes.output, 0)" }
         }
       ]
@@ -471,23 +251,21 @@
     {
       "id": "rank3_window_tiled_real",
       "priority": 19,
-      "requires": { "limits": { "maxComputeInvocationsPerWorkgroup": 256, "maxComputeWorkgroupSizeX": 256 } },
-      "tunables": { "WORKGROUP_SIZE": 256 },
       "when": ["rank3WindowLayout", "rank3WindowSpectrum", "rank3WindowTiledRealFits"],
       "passes": [
         {
           "id": "main",
           "name": "STFT.TiledReal",
-          "source": {
-            "shader": "dft-tiled-real.wgsl.jinja",
-            "inputs": {
-              "mode": "\"stft\"",
-              "transformLength": "dim(shapes.window, 0)",
-              "numBins": "dim(shapes.output, 2)",
-              "workgroupSize": "tunables.WORKGROUP_SIZE"
-            }
           },
-          "bindings": "stftWindowFft",
           "dispatch": {
             "x": "ceilDiv(dim(shapes.output, 2), 256)",
             "y": "dim(shapes.output, 1)",
@@ -499,28 +277,28 @@
     {
       "id": "rank3_window_rfft_packed",
       "priority": 22,
-      "tunables": { "WORKGROUP_SIZE": 64 },
       "when": ["rank3WindowLayout", "rank3WindowSpectrum", "rank3WindowSharedFftFits", "attrs.onesided != 0", "dim(shapes.signal, 2) == 1", "dim(shapes.window, 0) >= 64", "dim(shapes.output, 0) * dim(shapes.output, 1) >= tunables.TWIDDLE_TABLE_MIN_TRANSFORMS"],
       "intermediates": [{ "id": "fft_twiddles", "dtype": "float32", "shape": "[2 * dim(shapes.window, 0)]" }],
       "passes": [
         {
           "id": "twiddles",
           "name": "STFT.RealPackedTwiddleTable",
-          "source": {
-            "shader": "fft-stockham-twiddles.wgsl.jinja",
-            "inputs": { "fftLength": "dim(shapes.window, 0)", "realPacked": 1 }
-          },
-          "bindings": "stftTwiddleWrite",
-          "dispatch": { "threads": "dim(shapes.window, 0)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         },
         {
           "id": "main",
           "name": "STFT.RealPackedFftShared",
-          "source": {
-            "shader": "fft-real-packed-stockham-shared.wgsl.jinja",
-            "inputs": { "fftLength": "dim(shapes.window, 0)", "transformsPerWorkgroup": "packedTransformsPerWorkgroup" }
-          },
-          "bindings": "stftWindowFftTwiddles",
           "dispatch": {
             "x": "ceilDiv(dim(shapes.output, 1), packedTransformsPerWorkgroup)",
             "y": "dim(shapes.output, 0)"
@@ -537,28 +315,28 @@
         {
           "id": "twiddles",
           "name": "STFT.FftTwiddleTable",
-          "source": {
-            "shader": "fft-stockham-twiddles.wgsl.jinja",
-            "inputs": { "fftLength": "dim(shapes.window, 0)", "realPacked": 0 }
-          },
-          "bindings": "stftTwiddleWrite",
-          "dispatch": { "threads": "dim(shapes.window, 0) - 1", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         },
         {
           "id": "main",
           "name": "STFT.FftSharedTwiddles",
-          "source": {
-            "shader": "fft-stockham-shared.wgsl.jinja",
-            "inputs": {
-              "mode": "\"stft\"",
-              "fftLength": "dim(shapes.window, 0)",
-              "numBins": "dim(shapes.output, 2)",
-              "inComp": "dim(shapes.signal, 2)",
-              "inverse": false,
-              "useTwiddleTable": true
-            }
           },
-          "bindings": "stftWindowFftTwiddles",
           "dispatch": { "x": "dim(shapes.output, 1)", "y": "dim(shapes.output, 0)" }
         }
       ]
@@ -571,19 +349,17 @@
         {
           "id": "main",
           "name": "STFT.FftMixedRecurrent",
-          "source": {
-            "shader": "fft-stockham-shared.wgsl.jinja",
-            "inputs": {
-              "mode": "\"stft\"",
-              "fftLength": "dim(shapes.window, 0)",
-              "radix2Length": "rank3WindowFftRadix2",
-              "mixedRadixTail": true,
-              "numBins": "dim(shapes.output, 2)",
-              "inComp": "dim(shapes.signal, 2)",
-              "inverse": false
-            }
           },
-          "bindings": "stftWindowFft",
           "dispatch": { "x": "dim(shapes.output, 1)", "y": "dim(shapes.output, 0)" }
         }
       ]
@@ -596,17 +372,15 @@
         {
           "id": "main",
           "name": "STFT.FftShared",
-          "source": {
-            "shader": "fft-stockham-shared.wgsl.jinja",
-            "inputs": {
-              "mode": "\"stft\"",
-              "fftLength": "dim(shapes.window, 0)",
-              "numBins": "dim(shapes.output, 2)",
-              "inComp": "dim(shapes.signal, 2)",
-              "inverse": false
-            }
           },
-          "bindings": "stftWindowFft",
           "dispatch": { "x": "dim(shapes.output, 1)", "y": "dim(shapes.output, 0)" }
         }
       ]
@@ -614,15 +388,20 @@
     {
       "id": "rank3_window",
       "priority": 0,
-      "supersededBy": ["rank3_window_fft_shared_inplace", "rank3_window_fft_global_stages", "rank3_window_fft_storage_inplace", "rank3_window_fft_mixed", "rank3_window_rfft_packed", "rank3_window_fft_twiddles", "rank3_window_fft", "rank3_window_tiled_real"],
       "when": ["rank3WindowLayout", "rank3WindowSpectrum"],
       "passes": [
         {
           "id": "main",
           "name": "STFT",
-          "source": { "shader": "stft-general.wgsl.jinja", "inputs": { "hasWindowInput": true } },
-          "bindings": "stftWindowGeneral",
-          "dispatch": { "threads": "numel(shapes.output) / 2", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
@@ -634,9 +413,14 @@
         {
           "id": "main",
           "name": "STFT.GeneralNoWindow",
-          "source": { "shader": "stft-general.wgsl.jinja", "inputs": { "hasWindowInput": false } },
-          "bindings": "stftNoWindowGeneral",
-          "dispatch": { "threads": "numel(shapes.output) / 2", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     }

   "domain": "ai.onnx",
   "name": "STFT",
   "sinceVersion": 17,
+  "inputs": { "signal": { "dtype": "T", "rank": 3 }, "window": { "dtype": "T", "rank": 1, "optional": true } },
+  "outputs": {
+    "output": {
       "dtype": "T",
       "rank": 4,
       "shape": ["dim(shapes.signal, 0)", "expectedFrames", "floor(frameLength / 2) + 1 if attrs.onesided == 1 else frameLength", 2]
     }
   },
   "args": {
+    "frameStep": { "kind": "i32", "onnx": "frame_step" },
+    "frameLength": { "kind": "i32", "onnx": "frame_length", "required": false }
   },
+  "attributes": { "onesided": { "default": 1 } },
+  "attributeConstraints": { "onesided": { "values": [0, 1] } },
+  "typeConstraints": { "T": ["float32"] },
   "tunables": {
+    "WORKGROUP_SIZE": { "default": 128 },
+    "TWIDDLE_TABLE_MIN_TRANSFORMS": { "default": 4096 },
+    "TILED_REAL_MIN_LENGTH": { "default": 512 },
+    "GLOBAL_STAGE_LOCAL_LENGTH": { "default": 256 },
+    "GLOBAL_STAGE_MAX_DEVICE_STORAGE_BYTES": { "default": 16384 }
   },
   "derive": {
     "windowLength": "dim(shapes.window, 0) if (present.window and ranks.window == 1) else 0",
     "rank3OutputLayout": "ranks.signal == 3 and ranks.output == 4 and dim(shapes.output, 0) == dim(shapes.signal, 0) and dim(shapes.output, 1) == expectedFrames and dim(shapes.signal, 2) >= 1 and dim(shapes.signal, 2) <= 2 and dim(shapes.output, 3) == 2",
     "rank3WindowLayout": "present.window and ranks.window == 1 and rank3OutputLayout and frameControlsValid",
     "rank3WindowSpectrum": "false if not present.window else ((attrs.onesided == 0 and dim(shapes.output, 2) == dim(shapes.window, 0)) or (attrs.onesided != 0 and dim(shapes.signal, 2) == 1 and dim(shapes.output, 2) == floor(dim(shapes.window, 0) / 2) + 1))",
+    "rank3WindowDispatchFits": "dim(shapes.output, 1) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dim(shapes.output, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "rank3WindowSharedFftFits": "false if not present.window else (rank3WindowDispatchFits and pow2ceil(dim(shapes.window, 0)) == dim(shapes.window, 0) and dim(shapes.window, 0) >= 4 and dim(shapes.window, 0) * 16 <= device.limits.maxComputeWorkgroupStorageSize)",
     "rank3WindowFftRadix2": "1 if not present.window else pick([[dim(shapes.window, 0) % 256 == 0, 256], [dim(shapes.window, 0) % 128 == 0, 128], [dim(shapes.window, 0) % 64 == 0, 64], [dim(shapes.window, 0) % 32 == 0, 32], [dim(shapes.window, 0) % 16 == 0, 16], [dim(shapes.window, 0) % 8 == 0, 8], [dim(shapes.window, 0) % 4 == 0, 4]], 1)",
     "rank3WindowFftTailRadix": "1 if not present.window else dim(shapes.window, 0) / rank3WindowFftRadix2",
     "rank3WindowRadix2ScratchFits": "false if not present.window else (rank3WindowDispatchFits and pow2ceil(dim(shapes.window, 0)) == dim(shapes.window, 0) and dim(shapes.window, 0) >= 4 and dim(shapes.window, 0) * 16 > device.limits.maxComputeWorkgroupStorageSize and dim(shapes.output, 0) * dim(shapes.output, 1) > 0 and dim(shapes.output, 0) * dim(shapes.output, 1) * dim(shapes.window, 0) * 2 * 4 <= device.limits.maxStorageBufferBindingSize and dim(shapes.output, 0) * dim(shapes.output, 1) * dim(shapes.window, 0) * 2 * 4 <= device.limits.maxBufferSize)",
     "rank3WindowInplaceFftFits": "false if not present.window else (rank3WindowDispatchFits and pow2ceil(dim(shapes.window, 0)) == dim(shapes.window, 0) and dim(shapes.window, 0) >= 4 and dim(shapes.window, 0) * 16 > device.limits.maxComputeWorkgroupStorageSize and dim(shapes.window, 0) * 8 <= device.limits.maxComputeWorkgroupStorageSize and dim(shapes.output, 0) * dim(shapes.output, 1) > 0)",
     "rank3WindowGlobalFftFits": "false if not present.window else (rank3WindowRadix2ScratchFits and dim(shapes.window, 0) >= 2 * tunables.GLOBAL_STAGE_LOCAL_LENGTH and dim(shapes.window, 0) * 8 > device.limits.maxComputeWorkgroupStorageSize)",
+    "rank3WindowTiledRealFits": "false if not present.window else (rank3WindowDispatchFits and ranks.signal == 3 and dim(shapes.signal, 2) == 1 and dim(shapes.window, 0) >= tunables.TILED_REAL_MIN_LENGTH and dim(shapes.window, 0) * 4 <= device.limits.maxComputeWorkgroupStorageSize and dim(shapes.output, 0) > 0 and dim(shapes.output, 1) > 0 and ceilDiv(dim(shapes.output, 2), 256) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
     "packedFourWayFits": "false if not present.window else (device.limits.maxComputeInvocationsPerWorkgroup >= 256 and device.limits.maxComputeWorkgroupSizeX >= 256 and dim(shapes.window, 0) * 32 <= device.limits.maxComputeWorkgroupStorageSize)",
     "packedTransformsPerWorkgroup": "4 if packedFourWayFits else 1",
     "rank3NoWindowLayout": "not present.window and args.frameLength is defined and rank3OutputLayout and frameControlsValid",
     "rank3NoWindowSpectrum": "(attrs.onesided == 0 and dim(shapes.output, 2) == frameLength) or (attrs.onesided != 0 and dim(shapes.signal, 2) == 1 and dim(shapes.output, 2) == floor(frameLength / 2) + 1)",
+    "rank3NoWindowSharedFftFits": "false if ranks.signal != 3 else ((attrs.onesided == 0 or (dim(shapes.signal, 2) == 1 and dim(shapes.output, 2) >= 3)) and pow2ceil(frameLength) == frameLength and frameLength >= 4 and frameLength * 16 <= device.limits.maxComputeWorkgroupStorageSize and dim(shapes.output, 1) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and dim(shapes.output, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535))"
   },
+  "bindings": {
+    "signal": { "buffer": "read-only-storage", "elementType": "f32" },
+    "output": { "buffer": "storage", "elementType": "f32" },
+    "params": {
+      "buffer": "uniform",
+      "struct": [
+        { "name": "signalLength", "type": "u32", "value": "dim(shapes.signal, 1)" },
+        { "name": "numFrames", "type": "u32", "value": "dim(shapes.output, 1)" },
+        { "name": "frameStep", "type": "u32" }
+      ]
+    },
+    "window": { "buffer": "read-only-storage", "elementType": "f32" },
+    "scratch": { "scratch": "fft_scratch", "buffer": "storage", "elementType": "f32" },
+    "twiddles": { "scratch": "fft_twiddles", "buffer": "storage", "elementType": "f32" },
+    "twiddles_2": { "scratch": "fft_twiddles", "name": "twiddles", "buffer": "read-only-storage", "elementType": "f32" },
+    "params_2": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "count", "type": "u32", "value": "numel(shapes.output) / 2" },
+        { "name": "signalLength", "type": "u32", "value": "dim(shapes.signal, 1)" },
+        { "name": "signalComponents", "type": "u32", "value": "dim(shapes.signal, 2)" },
+        { "name": "frameLength", "type": "u32", "value": "dim(shapes.window, 0)" },
+        { "name": "frames", "type": "u32", "value": "dim(shapes.output, 1)" },
+        { "name": "bins", "type": "u32", "value": "dim(shapes.output, 2)" },
+        { "name": "frameStep", "type": "u32" }
+      ]
+    },
+    "params_3": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "count", "type": "u32", "value": "numel(shapes.output) / 2" },
+        { "name": "signalLength", "type": "u32", "value": "dim(shapes.signal, 1)" },
+        { "name": "signalComponents", "type": "u32", "value": "dim(shapes.signal, 2)" },
+        { "name": "frameLength", "type": "u32" },
+        { "name": "frames", "type": "u32", "value": "dim(shapes.output, 1)" },
+        { "name": "bins", "type": "u32", "value": "dim(shapes.output, 2)" },
+        { "name": "frameStep", "type": "u32" }
+      ]
+    }
   },
   "variants": [
     {
         {
           "id": "main",
           "name": "STFT.NoWindowFftShared",
+          "shader": "fft-stockham-shared.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"stft\"",
+            "fftLength": "frameLength",
+            "numBins": "dim(shapes.output, 2)",
+            "inComp": "dim(shapes.signal, 2)",
+            "inverse": false,
+            "hasWindow": false
           },
+          "bindings": ["signal", "output", "params"],
           "dispatch": { "x": "dim(shapes.output, 1)", "y": "dim(shapes.output, 0)" }
         }
       ]
     {
       "id": "rank3_window_fft_shared_inplace",
       "priority": 32,
       "when": ["rank3WindowLayout", "rank3WindowSpectrum", "rank3WindowInplaceFftFits"],
+      "requires": { "limits": { "maxComputeInvocationsPerWorkgroup": 256, "maxComputeWorkgroupSizeX": 256 } },
+      "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
       "passes": [
         {
           "id": "main",
           "name": "STFT.FftSharedInplace",
+          "shader": "fft-radix2-dit-shared-inplace.wgsl.jinja",
+          "derive": {
+            "fftLength": "dim(shapes.window, 0)",
+            "log2Length": "rank3WindowFftLog2",
+            "numBins": "dim(shapes.output, 2)",
+            "inComp": "dim(shapes.signal, 2)"
           },
+          "bindings": ["signal", "window", "output", "params"],
           "dispatch": { "x": "dim(shapes.output, 1)", "y": "dim(shapes.output, 0)" }
         }
       ]
     {
       "id": "rank3_window_fft_global_stages",
       "priority": 31,
       "when": ["rank3WindowLayout", "rank3WindowSpectrum", "rank3WindowGlobalFftFits"],
+      "demoteWhen": ["device.limits.maxComputeWorkgroupStorageSize > tunables.GLOBAL_STAGE_MAX_DEVICE_STORAGE_BYTES"],
+      "requires": { "limits": { "maxComputeInvocationsPerWorkgroup": 256, "maxComputeWorkgroupSizeX": 256 } },
+      "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
       "intermediates": [
         {
           "id": "fft_scratch",
         {
           "id": "local_prefix",
           "name": "STFT.FftLocalPrefix",
+          "shader": "fft-radix2-dit-storage-pass.wgsl.jinja",
+          "derive": {
+            "phase": "\"local_prefix\"",
+            "fftLength": "dim(shapes.window, 0)",
+            "log2Length": "rank3WindowFftLog2",
+            "localLength": "tunables.GLOBAL_STAGE_LOCAL_LENGTH",
+            "localLog2Length": "log2ceil(tunables.GLOBAL_STAGE_LOCAL_LENGTH)",
+            "numBins": "dim(shapes.output, 2)",
+            "inComp": "dim(shapes.signal, 2)"
           },
+          "bindings": ["signal", "window", "scratch", "params"],
           "dispatch": {
+            "x": "min(dim(shapes.output, 0) * dim(shapes.output, 1) * dim(shapes.window, 0) / tunables.GLOBAL_STAGE_LOCAL_LENGTH, 65535)",
+            "y": "ceilDiv(dim(shapes.output, 0) * dim(shapes.output, 1) * dim(shapes.window, 0) / tunables.GLOBAL_STAGE_LOCAL_LENGTH, 65535)",
+            "z": 1
           }
         },
         {
             {
               "id": "butterflies",
               "name": "STFT.FftGlobalStage",
+              "shader": "fft-radix2-dit-storage-pass.wgsl.jinja",
+              "derive": {
+                "phase": "\"stage\"",
+                "stage": "repeat.stage + log2ceil(tunables.GLOBAL_STAGE_LOCAL_LENGTH) + 1",
+                "fftLength": "dim(shapes.window, 0)",
+                "log2Length": "rank3WindowFftLog2",
+                "numBins": "dim(shapes.output, 2)",
+                "inComp": "dim(shapes.signal, 2)"
               },
+              "bindings": ["scratch"],
               "dispatch": {
+                "x": "min(ceilDiv((dim(shapes.output, 0) * dim(shapes.output, 1) * dim(shapes.window, 0) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
+                "y": "ceilDiv(ceilDiv((dim(shapes.output, 0) * dim(shapes.output, 1) * dim(shapes.window, 0) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
+                "z": 1
               }
             }
           ]
         {
           "id": "finish",
           "name": "STFT.FftGlobalFinish",
+          "shader": "fft-radix2-dit-storage-pass.wgsl.jinja",
+          "derive": {
+            "phase": "\"finish\"",
+            "fftLength": "dim(shapes.window, 0)",
+            "log2Length": "rank3WindowFftLog2",
+            "numBins": "dim(shapes.output, 2)",
+            "inComp": "dim(shapes.signal, 2)"
           },
+          "bindings": [
+            { "scratch": "fft_scratch", "name": "scratch", "buffer": "read-only-storage", "elementType": "f32" },
+            "output"
+          ],
           "dispatch": {
+            "x": "min(ceilDiv((dim(shapes.output, 0) * dim(shapes.output, 1) * dim(shapes.output, 2)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((dim(shapes.output, 0) * dim(shapes.output, 1) * dim(shapes.output, 2)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         }
       ]
         {
           "id": "main",
           "name": "STFT.FftStorageInplace",
+          "shader": "fft-radix2-dit-storage.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"stft\"",
+            "fftLength": "dim(shapes.window, 0)",
+            "log2Length": "rank3WindowFftLog2",
+            "numBins": "dim(shapes.output, 2)",
+            "inComp": "dim(shapes.signal, 2)",
+            "inverse": false
           },
+          "bindings": ["signal", "window", "scratch", "output", "params"],
           "dispatch": { "x": "dim(shapes.output, 1)", "y": "dim(shapes.output, 0)" }
         }
       ]
     {
       "id": "rank3_window_tiled_real",
       "priority": 19,
       "when": ["rank3WindowLayout", "rank3WindowSpectrum", "rank3WindowTiledRealFits"],
+      "requires": { "limits": { "maxComputeInvocationsPerWorkgroup": 256, "maxComputeWorkgroupSizeX": 256 } },
+      "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
       "passes": [
         {
           "id": "main",
           "name": "STFT.TiledReal",
+          "shader": "dft-tiled-real.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"stft\"",
+            "transformLength": "dim(shapes.window, 0)",
+            "numBins": "dim(shapes.output, 2)",
+            "workgroupSizeSpec": "tunables.WORKGROUP_SIZE"
           },
+          "bindings": ["signal", "window", "output", "params"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.output, 2), 256)",
             "y": "dim(shapes.output, 1)",
     {
       "id": "rank3_window_rfft_packed",
       "priority": 22,
       "when": ["rank3WindowLayout", "rank3WindowSpectrum", "rank3WindowSharedFftFits", "attrs.onesided != 0", "dim(shapes.signal, 2) == 1", "dim(shapes.window, 0) >= 64", "dim(shapes.output, 0) * dim(shapes.output, 1) >= tunables.TWIDDLE_TABLE_MIN_TRANSFORMS"],
+      "tunables": { "WORKGROUP_SIZE": { "default": 64 } },
       "intermediates": [{ "id": "fft_twiddles", "dtype": "float32", "shape": "[2 * dim(shapes.window, 0)]" }],
       "passes": [
         {
           "id": "twiddles",
           "name": "STFT.RealPackedTwiddleTable",
+          "shader": "fft-stockham-twiddles.wgsl.jinja",
+          "derive": { "fftLength": "dim(shapes.window, 0)", "realPacked": 1 },
+          "bindings": ["twiddles"],
+          "dispatch": {
+            "x": "min(ceilDiv((dim(shapes.window, 0)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((dim(shapes.window, 0)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         },
         {
           "id": "main",
           "name": "STFT.RealPackedFftShared",
+          "shader": "fft-real-packed-stockham-shared.wgsl.jinja",
+          "derive": { "fftLength": "dim(shapes.window, 0)", "transformsPerWorkgroup": "packedTransformsPerWorkgroup" },
+          "bindings": ["signal", "window", "twiddles_2", "output", "params"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.output, 1), packedTransformsPerWorkgroup)",
             "y": "dim(shapes.output, 0)"
         {
           "id": "twiddles",
           "name": "STFT.FftTwiddleTable",
+          "shader": "fft-stockham-twiddles.wgsl.jinja",
+          "derive": { "fftLength": "dim(shapes.window, 0)", "realPacked": 0 },
+          "bindings": ["twiddles"],
+          "dispatch": {
+            "x": "min(ceilDiv((dim(shapes.window, 0) - 1), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((dim(shapes.window, 0) - 1), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         },
         {
           "id": "main",
           "name": "STFT.FftSharedTwiddles",
+          "shader": "fft-stockham-shared.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"stft\"",
+            "fftLength": "dim(shapes.window, 0)",
+            "numBins": "dim(shapes.output, 2)",
+            "inComp": "dim(shapes.signal, 2)",
+            "inverse": false,
+            "useTwiddleTable": true
           },
+          "bindings": ["signal", "window", "twiddles_2", "output", "params"],
           "dispatch": { "x": "dim(shapes.output, 1)", "y": "dim(shapes.output, 0)" }
         }
       ]
         {
           "id": "main",
           "name": "STFT.FftMixedRecurrent",
+          "shader": "fft-stockham-shared.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"stft\"",
+            "fftLength": "dim(shapes.window, 0)",
+            "radix2Length": "rank3WindowFftRadix2",
+            "mixedRadixTail": true,
+            "numBins": "dim(shapes.output, 2)",
+            "inComp": "dim(shapes.signal, 2)",
+            "inverse": false
           },
+          "bindings": ["signal", "window", "output", "params"],
           "dispatch": { "x": "dim(shapes.output, 1)", "y": "dim(shapes.output, 0)" }
         }
       ]
         {
           "id": "main",
           "name": "STFT.FftShared",
+          "shader": "fft-stockham-shared.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"stft\"",
+            "fftLength": "dim(shapes.window, 0)",
+            "numBins": "dim(shapes.output, 2)",
+            "inComp": "dim(shapes.signal, 2)",
+            "inverse": false
           },
+          "bindings": ["signal", "window", "output", "params"],
           "dispatch": { "x": "dim(shapes.output, 1)", "y": "dim(shapes.output, 0)" }
         }
       ]
     {
       "id": "rank3_window",
       "priority": 0,
       "when": ["rank3WindowLayout", "rank3WindowSpectrum"],
+      "supersededBy": ["rank3_window_fft_shared_inplace", "rank3_window_fft_global_stages", "rank3_window_fft_storage_inplace", "rank3_window_fft_mixed", "rank3_window_rfft_packed", "rank3_window_fft_twiddles", "rank3_window_fft", "rank3_window_tiled_real"],
       "passes": [
         {
           "id": "main",
           "name": "STFT",
+          "shader": "stft-general.wgsl.jinja",
+          "derive": { "hasWindowInput": true },
+          "bindings": ["signal", "window", "output", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.output) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.output) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
         {
           "id": "main",
           "name": "STFT.GeneralNoWindow",
+          "shader": "stft-general.wgsl.jinja",
+          "derive": { "hasWindowInput": false },
+          "bindings": ["signal", "output", "params_3"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.output) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.output) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,25 +1,40 @@
 {
   "name": "ai.onnx.STFT",
-  "id": "_ai_onnx_stft_webgpu_285ab84",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "OKXSfrXKgDJBKuWa61pzqlE0q6XeDp3+Ac78kZJzvEs=",
-      "dft-tiled-real.wgsl.jinja": "jWXC4IENyJwpE7YFeiMBTtrmtncyfXddBe6JQBB3+h4=",
-      "fft-radix2-dit-shared-inplace.wgsl.jinja": "IGz/c29DBNlDXBpYQk9d44DHXtmORm86rJSyJxDQ4oU=",
-      "fft-radix2-dit-storage-pass.wgsl.jinja": "kDmHLr6NfPBnO2OPHYAkBgCA/dUzsHq9yIJS/v4eIpk=",
-      "fft-radix2-dit-storage.wgsl.jinja": "rlD1QCYyRVzQf4VEZTseGFYvUzjKNRneHUatxpz5Y9k=",
-      "fft-real-packed-stockham-shared.wgsl.jinja": "w9XapiIdKLjNy/3BwXrzmnBzUYWCzSvOoHE6lMQQJok=",
-      "fft-stockham-shared.wgsl.jinja": "ihNoaimYQY8F/oomzVbEh9g/g08yolKLYVQn79cZG1o=",
-      "fft-stockham-twiddles.wgsl.jinja": "6/+Ugc8XhOPl7iNTPc0Lwd0Moh8AFw+RVQZSWCxJbKY=",
-      "manifest.json": "stekrpAUb0OTNiTqln6FHfTBeHeZmISRKX/dP8eEkVU=",
-      "stft-general.wgsl.jinja": "UGbzUabSlzivMPxP2LJ2YRNK5Z6msagBugJvn3Tpu0s=",
-      "test.json": "Y05Ud3usGRQ2tsWy3twV3mndN+3FjHo3LjOwTA24Nog="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.STFT" }
 }

 {
   "name": "ai.onnx.STFT",
+  "id": "_ai_onnx_stft_webgpu_06cb038",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "vHcB5tti+CM3DMtdpAb1k6Lb72nQ3dLXXqVWPYHWi+Y=",
+      "dft-tiled-real.wgsl.jinja": "nls3dnjWMqMRdmwudzIVekmyeWr13k1DaSsSxfdX+Xo=",
+      "fft-radix2-dit-shared-inplace.wgsl.jinja": "uKO/mdFdypgtav2qghpelG0iwgR4WXM/mpgJVy2+MQA=",
+      "fft-radix2-dit-storage-pass.wgsl.jinja": "lSFcs2O2JtHaV8WGmyvtfo/QZi6xaflVR71cOip6zJ8=",
+      "fft-radix2-dit-storage.wgsl.jinja": "Rc4YVTVGib28SSfikk4IPyLTDL3BlEGCAsH04WyfqbE=",
+      "fft-real-packed-stockham-shared.wgsl.jinja": "fPcCFvTUojCPabsRnC9y7GvJVkogRQ7zcQtC48YAR4s=",
+      "fft-stockham-shared.wgsl.jinja": "W1qZ/FnWKVSZSpE43qo/06isxY6D+Va6+tzG7/myock=",
+      "fft-stockham-twiddles.wgsl.jinja": "SVeaa/kM8Fx9RHH3yVcvx/twJqPjZtrffvi0pGesdyg=",
+      "manifest.json": "8J0Y8qFjzAsLgbgEydZN8iD67cDTFYyOs1daSYYUeaA=",
+      "stft-general.wgsl.jinja": "R006F8Q2a4U7S+V1FyFsc+nUUMjx5ya4yGOFkcIiW+w=",
+      "test.json": "R2V8u8+fDmmwimCB8Cwj+xradfacxn/hrIJ2L/oF49Q="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "rank3_no_window_fft": ["fft-stockham-shared.wgsl.jinja"],
+      "rank3_window_fft_shared_inplace": ["fft-radix2-dit-shared-inplace.wgsl.jinja"],
+      "rank3_window_fft_global_stages": ["fft-radix2-dit-storage-pass.wgsl.jinja"],
+      "rank3_window_fft_storage_inplace": ["fft-radix2-dit-storage.wgsl.jinja"],
+      "rank3_window_tiled_real": ["dft-tiled-real.wgsl.jinja"],
+      "rank3_window_rfft_packed": ["fft-real-packed-stockham-shared.wgsl.jinja", "fft-stockham-twiddles.wgsl.jinja"],
+      "rank3_window_fft_twiddles": ["fft-stockham-shared.wgsl.jinja", "fft-stockham-twiddles.wgsl.jinja"],
+      "rank3_window_fft_mixed": ["fft-stockham-shared.wgsl.jinja"],
+      "rank3_window_fft": ["fft-stockham-shared.wgsl.jinja"],
+      "rank3_window": ["stft-general.wgsl.jinja"],
+      "rank3_no_window": ["stft-general.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/stft-general.wgsl.jinja CHANGED Viewed

@@ -1,30 +1,30 @@
 {% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
 {% if note == "dispatch-limit" %}
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
 {% elif note == "limit" %}
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit.
 {% elif note == "device-axis" %}
-  // The flat dispatch is folded across x/y at the device's per-axis workgroup
-  // limit; gid.y carries the high portion of the output index.
 {% elif note == "vec4-limit" %}
-  // 2D-folded flat vec4 index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit (the dispatch caps x and spills into y).
 {% elif note == "element-limit" %}
   // 2D-folded flat element index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit.
 {% elif note == "dispatch" %}
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension dispatch limit.
 {% endif %}
 {% if bound == "" %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
 {%- elif guardInline %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) { return; }
 {%- else %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) {
     return;
   }
@@ -33,18 +33,17 @@
 {{ env.wgsl.resourceDeclarations }}
 // Backend-stable unit-circle twiddles for DFT kernels. Retain the phase as a
-// rational turn and route all polynomial evaluation through the shared precise
-// trig core so FFT, direct DFT, STFT, and cosine windows agree exactly.
 //
 // phase is in [0, period). sign is -1 for a forward transform and +1 for an
 // inverse transform. The result is (cos, sin).
 // Backend-stable f32 sine/cosine core.
 //
-// Shader transcendental accuracy is implementation-defined, and some portable
-// backends are only accurate to roughly 1e-4. Each path retains the most accurate
-// available phase representation, reduces it to [-pi, pi], then uses these
-// polynomials. The half-pi core is degree 13 for sine and degree 12 for cosine;
-// truncation error is well below one f32 ULP over its documented interval.
 const PRECISE_TRIG_PI: f32 = 3.141592653589793;
 const PRECISE_TRIG_TWO_PI: f32 = 6.283185307179586;
 const PRECISE_TRIG_HALF_PI: f32 = 1.5707963267948966;
@@ -118,7 +117,7 @@ fn dft_precise_twiddle(phase: u32, period: u32, sign: f32) -> vec2<f32> {
 // scalar. Window samples default to 1.0 when no window input is present.
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
-fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
 {{ flat_index_2d("index", guardInline=true, note="") }}
   let bin = index % params.bins;
@@ -136,7 +135,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
   for (var n = 0u; n < frame_length; n = n + 1u) {
     let signal_index = start + n;
     let signal_base = (batch * params.signalLength + signal_index) * params.signalComponents;
-{% if source.hasWindowInput %}
     let w = window[n];
 {% else %}
     let w = 1.0;

 {% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
 {% if note == "dispatch-limit" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width (outputs > 16.7M elements).
 {% elif note == "limit" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
 {% elif note == "device-axis" %}
+  // The flat dispatch is folded across x/y at a fixed per-axis workgroup
+  // width; gid.y carries the high portion of the output index.
 {% elif note == "vec4-limit" %}
+  // 2D-folded flat vec4 index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width (the dispatch caps x and spills into y).
 {% elif note == "element-limit" %}
   // 2D-folded flat element index: gid.y carries the high bits past the
+  // dispatch's per-axis workgroup fold width.
 {% elif note == "dispatch" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
 {% endif %}
 {% if bound == "" %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
 {%- elif guardInline %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) { return; }
 {%- else %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) {
     return;
   }
 {{ env.wgsl.resourceDeclarations }}
 // Backend-stable unit-circle twiddles for DFT kernels. Retain the phase as a
+// rational turn and route polynomial evaluation through one trigonometric core
+// for FFT, direct DFT, STFT, and cosine-window calculations.
 //
 // phase is in [0, period). sign is -1 for a forward transform and +1 for an
 // inverse transform. The result is (cos, sin).
 // Backend-stable f32 sine/cosine core.
 //
+// Shader transcendental accuracy is implementation-defined. Each path retains
+// its available phase representation, reduces it to [-pi, pi], then evaluates
+// explicit polynomials in a fixed order. The half-pi core is degree 13 for sine
+// and degree 12 for cosine.
 const PRECISE_TRIG_PI: f32 = 3.141592653589793;
 const PRECISE_TRIG_TWO_PI: f32 = 6.283185307179586;
 const PRECISE_TRIG_HALF_PI: f32 = 1.5707963267948966;
 // scalar. Window samples default to 1.0 when no window input is present.
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 {{ flat_index_2d("index", guardInline=true, note="") }}
   let bin = index % params.bins;
   for (var n = 0u; n < frame_length; n = n + 1u) {
     let signal_index = start + n;
     let signal_base = (batch * params.signalLength + signal_index) * params.signalComponents;
+{% if hasWindowInput %}
     let w = window[n];
 {% else %}
     let w = 1.0;

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.STFT",
   "fixtureArrays": {
     "onnx_backend_arange128_cos_window16_input_window": [1, 0.961939766, 0.853553391, 0.691341716, 0.5, 0.308658284, 0.146446609, 0.038060234, 0, 0.038060234, 0.146446609, 0.308658284, 0.5, 0.691341716, 0.853553391, 0.961939766],
     "ort_complex_input_batched_twosided_frame_length_input_signal": [1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 99, 0, 99, 0, 99, 0, 99, 0, 99, 0, 99, 0, 99, 0, 99, 0]
@@ -18,7 +17,7 @@
     {
       "name": "many_frames_fft_twiddle_table_win64",
       "provenance": {
-        "notes": "Exercises the shared twiddle-table FFT selected for thousands of independent transforms without making correctness coverage depend on the much larger dispatch-cliff benchmark allocation."
       },
       "inputs": {
         "signal": { "dtype": "float32", "shape": [1, 4159, 1] },
@@ -30,7 +29,7 @@
     {
       "name": "many_frames_fft_twiddle_table_win64_twosided",
       "provenance": {
-        "notes": "Twosided sibling of the twiddle-table selector above. The packed real-FFT kernel that otherwise wins at this transform count requires onesided != 0, so a full-spectrum output is what leaves the shared twiddle-table FFT as the selected route."
       },
       "attrs": { "onesided": 0 },
       "inputs": {
@@ -53,7 +52,7 @@
     {
       "name": "many_frames_packed_rfft_win512_radix4",
       "provenance": {
-        "notes": "Selector and correctness lock for the packed real N=512 path. The 4097 transforms cross the twiddle-table threshold and leave one partial four-transform workgroup on capable devices."
       },
       "inputs": {
         "signal": {
@@ -425,7 +424,7 @@
     {
       "name": "fft_window512_real_twosided_step256",
       "provenance": {
-        "notes": "Real/full-spectrum coverage for the power-of-two FFT path; complements complex/twosided N=512 and real/onesided N=1024 fixtures."
       },
       "attrs": { "onesided": 0 },
       "inputs": {
@@ -489,7 +488,7 @@
     {
       "name": "fft_window2048_radix2_dit_onesided_min_storage",
       "provenance": {
-        "notes": "The in-place radix-2 rung is selected inside a BAND — rank3WindowInplaceFftFits wants window*16 > maxComputeWorkgroupStorageSize (the two-buffer form does not fit) and window*8 <= it (the in-place one does) — so the band moves with the device and no single window covers every tier. window 4096 is the band at 32768 (the adapter here); 2048 is the band at the 16384 WebGPU guaranteed minimum, which is what the min-limits tier reports. Without this sibling the rung is rendered on every low-storage tier and executed on none of them. Deliberately no requires.limits: outside its band the case is still valid, it just selects the staged rung."
       },
       "inputs": {
         "signal": {
@@ -684,7 +683,7 @@
     {
       "name": "mixed_stockham_n1000_real_twosided_tail125",
       "provenance": {
-        "notes": "Window length 1000 factors as 8 x 125 and exercises the shared Stockham mixed-radix route after the tail-bound transition. Odd lengths 513 and 777 retain independent tiled-real coverage."
       },
       "attrs": { "onesided": 0 },
       "inputs": {

 {
   "fixtureArrays": {
     "onnx_backend_arange128_cos_window16_input_window": [1, 0.961939766, 0.853553391, 0.691341716, 0.5, 0.308658284, 0.146446609, 0.038060234, 0, 0.038060234, 0.146446609, 0.308658284, 0.5, 0.691341716, 0.853553391, 0.961939766],
     "ort_complex_input_batched_twosided_frame_length_input_signal": [1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 99, 0, 99, 0, 99, 0, 99, 0, 99, 0, 99, 0, 99, 0, 99, 0]
     {
       "name": "many_frames_fft_twiddle_table_win64",
       "provenance": {
+        "notes": "A length-64 window over 4,159 samples produces 4,096 independent transforms and selects the shared twiddle-table FFT geometry."
       },
       "inputs": {
         "signal": { "dtype": "float32", "shape": [1, 4159, 1] },
     {
       "name": "many_frames_fft_twiddle_table_win64_twosided",
       "provenance": {
+        "notes": "A 4,096-frame, 64-point transform with `onesided=0` requires the full spectrum and exercises the shared twiddle-table FFT."
       },
       "attrs": { "onesided": 0 },
       "inputs": {
     {
       "name": "many_frames_packed_rfft_win512_radix4",
       "provenance": {
+        "notes": "A real 512-point STFT over 4,097 transforms exercises packed processing and leaves one partial four-transform workgroup when that specialization is available."
       },
       "inputs": {
         "signal": {
     {
       "name": "fft_window512_real_twosided_step256",
       "provenance": {
+        "notes": "A real length-512 transform with onesided=0 emits the full 512-bin spectrum; frame step 256 produces three frames."
       },
       "attrs": { "onesided": 0 },
       "inputs": {
     {
       "name": "fft_window2048_radix2_dit_onesided_min_storage",
       "provenance": {
+        "notes": "A 2048-point window exercises the in-place radix-2 specialization at the 16 KiB WebGPU workgroup-storage minimum: its 8 bytes per element scratch fits while the 16 bytes per element two-buffer scratch does not. On larger storage limits the case remains valid and may use the two-buffer specialization instead, so it requires no explicit device limit."
       },
       "inputs": {
         "signal": {
     {
       "name": "mixed_stockham_n1000_real_twosided_tail125",
       "provenance": {
+        "notes": "Window length 1,000 factors as 8 x 125, selecting the shared mixed-radix transform. With onesided=0, each frame emits the full 1,000-bin spectrum."
       },
       "attrs": { "onesided": 0 },
       "inputs": {