sync 91d990483a17

Browse files

Files changed (12) hide show

README.md +17 -13
build/webgpu/bench.json +0 -1
build/webgpu/dft-contiguous-naive.wgsl.jinja +11 -12
build/webgpu/dft-general.wgsl.jinja +24 -15
build/webgpu/dft-rank4.wgsl.jinja +21 -22
build/webgpu/dft-runtime-axis-fft-shared.wgsl.jinja +186 -0
build/webgpu/dft-tiled-real.wgsl.jinja +9 -10
build/webgpu/fft-radix2-dit-storage.wgsl.jinja +19 -20
build/webgpu/fft-stockham-shared.wgsl.jinja +30 -42
build/webgpu/manifest.json +393 -610
build/webgpu/metadata.json +31 -12
build/webgpu/test.json +115 -13

README.md CHANGED Viewed

@@ -18,17 +18,17 @@ See the [ONNX `DFT` spec](https://onnx.ai/onnx/operators/onnx__DFT.html) for the
 ## Inputs
-| Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
-| --- | --- | --- | --- | --- | --- | --- | --- |
-| `input` | `input` | `T` | same as logical dtype | — | — | Input signal tensor; the last dimension is 1 for real values or 2 for complex `(real, imaginary)` pairs. | required |
-| `dft_length` | `dft_length` | `L` | same as logical dtype | `0` | — | Optional int32 scalar controlling the signal length used for the transform; input is zero-padded or truncated to this length. | optional |
-| `axis` | `axis` | `I` | `int32` | `0` | — | Optional logical int64 scalar specifying the dimension over which to compute the DFT (the last axis is reserved for the real/imaginary component and is excluded); signed axes use int32 WebGPU storage and default to `-2` when omitted. | optional |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
-| --- | --- | --- | --- | --- | --- | --- |
-| `output` | `output` | `T` | same as `input` | — | DFT result tensor; last dimension is 2 (complex) for forward DFT and RFFT, or 1 (real) for IRFFT. | required |
 ## Attributes
@@ -49,28 +49,32 @@ 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
 - [`dft-contiguous-naive.wgsl.jinja`](build/webgpu/dft-contiguous-naive.wgsl.jinja)
 - [`dft-general.wgsl.jinja`](build/webgpu/dft-general.wgsl.jinja)
 - [`dft-rank4.wgsl.jinja`](build/webgpu/dft-rank4.wgsl.jinja)
 - [`dft-tiled-real.wgsl.jinja`](build/webgpu/dft-tiled-real.wgsl.jinja)
 - [`fft-radix2-dit-storage.wgsl.jinja`](build/webgpu/fft-radix2-dit-storage.wgsl.jinja)
 - [`fft-stockham-shared.wgsl.jinja`](build/webgpu/fft-stockham-shared.wgsl.jinja)
 ## Use with `@huggingface/kernels`
-The loader automatically allocates outputs whose metadata it can derive from the manifest contract and this call.
-The explicit `outputs` entries provide shape and logical dtype metadata for the results listed below:
-- `output`
-Each entry either requests an optional result or supplies metadata that cannot be inferred from the inputs.
 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 | WebGPU storage | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `input` | `T` | same as logical dtype | — | — | Input signal tensor; the last dimension is 1 for real values or 2 for complex `(real, imaginary)` pairs. | required |
+| `dft_length` | `L` | same as logical dtype | `0` | — | Optional int32 scalar controlling the signal length used for the transform; input is zero-padded or truncated to this length. | optional |
+| `axis` | `I` | `int32` | `0` | — | Optional logical int64 scalar specifying the dimension over which to compute the DFT (the last axis is reserved for the real/imaginary component and is excluded); signed axes use int32 WebGPU storage and default to `-2` when omitted. | optional |
 ## Outputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `output` | `T` | same as `input` | — | DFT result tensor; last dimension is 2 (complex) for forward DFT and RFFT, or 1 (real) for IRFFT. | required |
 ## Attributes
 ## 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
 - [`dft-contiguous-naive.wgsl.jinja`](build/webgpu/dft-contiguous-naive.wgsl.jinja)
 - [`dft-general.wgsl.jinja`](build/webgpu/dft-general.wgsl.jinja)
 - [`dft-rank4.wgsl.jinja`](build/webgpu/dft-rank4.wgsl.jinja)
+- [`dft-runtime-axis-fft-shared.wgsl.jinja`](build/webgpu/dft-runtime-axis-fft-shared.wgsl.jinja)
 - [`dft-tiled-real.wgsl.jinja`](build/webgpu/dft-tiled-real.wgsl.jinja)
 - [`fft-radix2-dit-storage.wgsl.jinja`](build/webgpu/fft-radix2-dit-storage.wgsl.jinja)
 - [`fft-stockham-shared.wgsl.jinja`](build/webgpu/fft-stockham-shared.wgsl.jinja)
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Outputs with inferable metadata are allocated automatically. Explicit `outputs` entries request optional results or provide metadata that cannot be inferred from the supplied inputs and attributes.
+This example supplies explicit metadata for:
+- `output`
 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.DFT",
   "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
   "cases": [
     {

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

build/webgpu/dft-contiguous-naive.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;
@@ -137,8 +136,8 @@ fn is_nyquist(freq: u32, signal_length: u32) -> bool {
 {% endmacro %}
-fn transform_input_base(sample: u32{% if source.batched %}, batch: u32{% endif %}) -> u32 {
-{% if source.batched %}
   return (batch * params.signalLength + sample) * params.inputComponents;
 {% else %}
   return sample * params.inputComponents;
@@ -146,14 +145,14 @@ fn transform_input_base(sample: u32{% if source.batched %}, batch: u32{% endif %
 }
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
-fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
   // Fold a long 1D output grid into y at the device dispatch limit.
-  let index = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
   if (index >= params.count) {
     return;
   }
-{% if source.batched %}
   let transform_index = index % params.outputLength;
   let batch = index / params.outputLength;
 {{ dft_naive_accumulate("transform_index", "params.signalLength", "params.outputLength", "batch", "0u", true, false) }}

 {{ 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;
 {% endmacro %}
+fn transform_input_base(sample: u32{% if batched %}, batch: u32{% endif %}) -> u32 {
+{% if batched %}
   return (batch * params.signalLength + sample) * params.inputComponents;
 {% else %}
   return sample * params.inputComponents;
 }
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   // Fold a long 1D output grid into y at the device dispatch limit.
+  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if (index >= params.count) {
     return;
   }
+{% if batched %}
   let transform_index = index % params.outputLength;
   let batch = index / params.outputLength;
 {{ dft_naive_accumulate("transform_index", "params.signalLength", "params.outputLength", "batch", "0u", true, false) }}

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

@@ -1,17 +1,19 @@
 {{ 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;
@@ -101,7 +103,7 @@ fn is_nyquist(freq: u32, signal_length: u32) -> bool {
 }
 fn resolve_axis() -> u32 {
-{% if source.hasAxisInput %}
   let raw = axis_in[0];
   let rank = i32(params.rank);
   let norm = select(raw, raw + rank, raw < 0);
@@ -112,9 +114,9 @@ fn resolve_axis() -> u32 {
 }
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
-fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
   // 2D-folded flat index over output elements (one per complex sample).
-  let index = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
   if (index >= params.count) { return; }
   let in_dims = array<u32, 7>(params.inD0, params.inD1, params.inD2, params.inD3, params.inD4, params.inD5, params.inD6);
@@ -124,13 +126,20 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
   let n_logical = params.rank - 1u;
   let pad = 7u - n_logical;
   let resolved_axis = resolve_axis();
-{% if source.skipRuntimeAxis1 %}
-  // Complementary to the guarded shared-memory FFT pass in the dynamic-axis
-  // specialization. Axis 1 is already complete, so this dispatch is a uniform
-  // no-op for that case.
   if (resolved_axis == 1u) { return; }
 {% endif %}
   let transform_slot = resolved_axis + pad;
   // Decode the right-aligned logical output coords from the flat index.
   var coords = array<u32, 7>(0u, 0u, 0u, 0u, 0u, 0u, 0u);
@@ -148,7 +157,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
   let inverse = params.inverse != 0u;
   let onesided_inverse = inverse && params.onesided != 0u;
   var signal_length = input_axis_length;
-{% if source.hasDftLengthInput %}
   if (dft_len_in[0] > 0) { signal_length = u32(dft_len_in[0]); }
 {% else %}
   if (params.dftLengthAttr > 0u) { signal_length = params.dftLengthAttr; }

+{% set skipSharedFftAxes = skipSharedFftAxes | default(false) %}
+{% set sharedFftMinLength = sharedFftMinLength | default(0) %}
+{% set sharedFftMaxLength = sharedFftMaxLength | default(0) %}
 {{ 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;
 }
 fn resolve_axis() -> u32 {
+{% if hasAxisInput %}
   let raw = axis_in[0];
   let rank = i32(params.rank);
   let norm = select(raw, raw + rank, raw < 0);
 }
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   // 2D-folded flat index over output elements (one per complex sample).
+  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if (index >= params.count) { return; }
   let in_dims = array<u32, 7>(params.inD0, params.inD1, params.inD2, params.inD3, params.inD4, params.inD5, params.inD6);
   let n_logical = params.rank - 1u;
   let pad = 7u - n_logical;
   let resolved_axis = resolve_axis();
+{% if skipRuntimeAxis1 %}
+  // In this specialization, a resolved axis of 1 returns uniformly without
+  // writing output elements.
   if (resolved_axis == 1u) { return; }
 {% endif %}
   let transform_slot = resolved_axis + pad;
+{% if skipSharedFftAxes %}
+  // Power-of-two axis lengths in the configured inclusive range return
+  // uniformly without writing output elements in this pass.
+  {
+    let axis_len = in_dims[transform_slot];
+    if (axis_len >= {{ sharedFftMinLength }}u && axis_len <= {{ sharedFftMaxLength }}u && (axis_len & (axis_len - 1u)) == 0u) { return; }
+  }
+{% endif %}
   // Decode the right-aligned logical output coords from the flat index.
   var coords = array<u32, 7>(0u, 0u, 0u, 0u, 0u, 0u, 0u);
   let inverse = params.inverse != 0u;
   let onesided_inverse = inverse && params.onesided != 0u;
   var signal_length = input_axis_length;
+{% if hasDftLengthInput %}
   if (dft_len_in[0] > 0) { signal_length = u32(dft_len_in[0]); }
 {% else %}
   if (params.dftLengthAttr > 0u) { signal_length = params.dftLengthAttr; }

build/webgpu/dft-rank4.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;
@@ -183,7 +182,7 @@ fn transform_input_base(sample: u32, batch: u32, other: u32) -> u32 {
 }
 @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) }}
   let out_d2 = index % params.outputD2;

 {% 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;
 }
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 {{ flat_index_2d("index", guardInline=true) }}
   let out_d2 = index % params.outputD2;

build/webgpu/dft-runtime-axis-fft-shared.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,186 @@

+{{ 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;
+// Input must be in [-pi/2, pi/2]. Returns (cos(x), sin(x)).
+fn precise_sincos_half_pi(x: f32) -> vec2<f32> {
+  let x2 = x * x;
+  var sinPolynomial = 1.6059043836821613e-10;
+  sinPolynomial = fma(sinPolynomial, x2, -2.505210838544172e-8);
+  sinPolynomial = fma(sinPolynomial, x2, 2.7557319223985893e-6);
+  sinPolynomial = fma(sinPolynomial, x2, -1.984126984126984e-4);
+  sinPolynomial = fma(sinPolynomial, x2, 8.333333333333333e-3);
+  sinPolynomial = fma(sinPolynomial, x2, -1.6666666666666666e-1);
+  let sine = x * fma(sinPolynomial, x2, 1.0);
+  var cosPolynomial = 2.08767569878681e-9;
+  cosPolynomial = fma(cosPolynomial, x2, -2.755731922398589e-7);
+  cosPolynomial = fma(cosPolynomial, x2, 2.48015873015873e-5);
+  cosPolynomial = fma(cosPolynomial, x2, -1.388888888888889e-3);
+  cosPolynomial = fma(cosPolynomial, x2, 4.1666666666666664e-2);
+  cosPolynomial = fma(cosPolynomial, x2, -5.0e-1);
+  let cosine = fma(cosPolynomial, x2, 1.0);
+  return vec2<f32>(cosine, sine);
+}
+// Input must be in [-pi, pi]. Returns (cos(x), sin(x)).
+fn precise_sincos_centered(x: f32) -> vec2<f32> {
+  var folded = x;
+  var cosineSign = 1.0;
+  if (folded > PRECISE_TRIG_HALF_PI) {
+    folded = PRECISE_TRIG_PI - folded;
+    cosineSign = -1.0;
+  } else if (folded < -PRECISE_TRIG_HALF_PI) {
+    folded = -PRECISE_TRIG_PI - folded;
+    cosineSign = -1.0;
+  }
+  let value = precise_sincos_half_pi(folded);
+  return vec2<f32>(cosineSign * value.x, value.y);
+}
+fn dft_precise_twiddle(phase: u32, period: u32, sign: f32) -> vec2<f32> {
+  let reducedPhase = phase % period;
+  if (reducedPhase == 0u) {
+    return vec2<f32>(1.0, sign * 0.0);
+  }
+  if (period % 4u == 0u) {
+    let quarter = period / 4u;
+    if (reducedPhase == quarter) {
+      return vec2<f32>(0.0, sign);
+    }
+    if (reducedPhase == 2u * quarter) {
+      return vec2<f32>(-1.0, sign * 0.0);
+    }
+    if (reducedPhase == 3u * quarter) {
+      return vec2<f32>(0.0, -sign);
+    }
+  }
+  // Center the rational turn before multiplying by 2*pi. This avoids the
+  // cancellation in `2*pi*(N-1)/N - 2*pi` for bins near the end of a row.
+  let turns = f32(reducedPhase) / f32(period);
+  let centeredTurns = turns - floor(turns + 0.5);
+  return precise_sincos_centered(sign * PRECISE_TRIG_TWO_PI * centeredTurns);
+}
+// Shared-memory radix-2 Stockham FFT along a dynamically selected transform
+// axis for inputs of any supported rank. The opset-20 `axis` input is a tensor,
+// so the transform axis, its length, and the (outer, inner) geometry are resolved
+// per dispatch from the logical dimensions in `params`. Shared scratch is sized
+// for `MAX_N`.
+// A resolved axis outside the inclusive `[MIN_N, MAX_N]` power-of-two range
+// returns uniformly before any barrier and does not write output elements.
+//
+// One transform per workgroup: dispatch covers the largest transform count any
+// admissible axis can need; workgroups past the resolved count exit uniformly.
+// Stage s of the log2(N) ping-pong reads a = src[j], b = src[j + N/2], applies
+// w = exp(SIGN*i*pi*(j mod ns)/ns) to b and writes the autosort pair, so the
+// natural-order spectrum lands in the last stage's destination half. The
+// twiddle angle never exceeds pi.
+const MIN_N: u32 = {{ fftMinLength }}u;
+const MAX_N: u32 = {{ fftMaxLength }}u;
+const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+const SIGN: f32 = {{ "1.0" if inverse else "-1.0" }};
+var<workgroup> shRe: array<f32, 2u * MAX_N>;
+var<workgroup> shIm: array<f32, 2u * MAX_N>;
+fn resolve_axis() -> u32 {
+  let raw = axis_in[0];
+  let rank = i32(params.rank);
+  let norm = select(raw, raw + rank, raw < 0);
+  return u32(norm);
+}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(workgroup_id) wgid: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>) {
+  let tid = lid.x;
+  let dims = array<u32, 7>(params.inD0, params.inD1, params.inD2, params.inD3, params.inD4, params.inD5, params.inD6);
+  let pad = 7u - (params.rank - 1u);
+  let slot = resolve_axis() + pad;
+  let n = dims[slot];
+  // This pass accepts power-of-two lengths in the inclusive configured range.
+  let admissible = n >= MIN_N && n <= MAX_N && (n & (n - 1u)) == 0u;
+  if (!admissible) { return; }
+  var outer = 1u;
+  var inner = 1u;
+  for (var s = 0u; s < 7u; s = s + 1u) {
+    let d = dims[s];
+    outer = select(outer, outer * d, s < slot);
+    inner = select(inner, inner * d, s > slot);
+  }
+  // Flat transform index over the 2D-folded workgroup grid (baked fold width).
+  let flat = wgid.x + wgid.y * {{ DISPATCH_FOLD_WIDTH }}u;
+  if (flat >= outer * inner) { return; }
+  let outerIndex = flat / inner;
+  let innerIndex = flat - outerIndex * inner;
+  for (var k = tid; k < n; k = k + WG) {
+    let base = ((outerIndex * n + k) * inner + innerIndex) * {{ inComp }}u;
+    shRe[k] = input[base];
+{% if inComp == 2 %}
+    shIm[k] = input[base + 1u];
+{% else %}
+    shIm[k] = 0.0;
+{% endif %}
+  }
+  workgroupBarrier();
+  let nHalf = n / 2u;
+  var srcOff = 0u;
+  var dstOff = MAX_N;
+  for (var s = 0u; (1u << s) < n; s = s + 1u) {
+    let ns = 1u << s;
+    for (var j = tid; j < nHalf; j = j + WG) {
+      let jm = j & (ns - 1u);
+      let outBase = ((j >> s) << (s + 1u)) + jm;
+      let ar = shRe[srcOff + j];
+      let ai = shIm[srcOff + j];
+      let br = shRe[srcOff + nHalf + j];
+      let bi = shIm[srcOff + nHalf + j];
+      let twiddle = dft_precise_twiddle(jm, 2u * ns, SIGN);
+      let c = twiddle.x;
+      let sn = twiddle.y;
+      let tr = br * c - bi * sn;
+      let ti = br * sn + bi * c;
+      shRe[dstOff + outBase] = ar + tr;
+      shIm[dstOff + outBase] = ai + ti;
+      shRe[dstOff + outBase + ns] = ar - tr;
+      shIm[dstOff + outBase + ns] = ai - ti;
+    }
+    workgroupBarrier();
+    let swap = srcOff;
+    srcOff = dstOff;
+    dstOff = swap;
+  }
+  // The final swap leaves the natural-order result at srcOff. In inverse mode,
+  // the scale is 1/n; this route selects a power-of-two n.
+{% if inverse %}
+  let scale = 1.0 / f32(n);
+{% else %}
+  let scale = 1.0;
+{% endif %}
+  for (var k = tid; k < n; k = k + WG) {
+    let outBase = ((outerIndex * n + k) * inner + innerIndex) * 2u;
+    output[outBase] = shRe[srcOff + k] * scale;
+    output[outBase + 1u] = shIm[srcOff + k] * scale;
+  }
+}

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-storage.wgsl.jinja CHANGED Viewed

@@ -8,12 +8,12 @@
 // 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" %}
-{% set USE_SCRATCH = source.mode != "dft" %}
 {% if not STFT %}
 // 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).
@@ -21,11 +21,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 %}
@@ -106,14 +105,14 @@ fn dft_precise_twiddle(phase: u32, period: u32, sign: f32) -> vec2<f32> {
 {% if STFT %}
 const PI: f32 = 3.141592653589793;
 {% endif %}
-const N: u32 = {{ source.fftLength }}u;
 const N_HALF: u32 = N / 2u;
-const LOG2_N: u32 = {{ source.log2Length }}u;
 {% if USE_SCRATCH %}
-const NUM_BINS: u32 = {{ source.numBins }}u;
 {% endif %}
 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(
@@ -138,16 +137,16 @@ fn main(
     var value = vec2<f32>(0.0);
 {% if STFT %}
     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 %}
 {% else %}
-    let inputBase = ((outerIndex * N + n) * params.inner + innerIndex) * {{ source.inComp }}u;
     value.x = input[inputBase];
-{% if source.inComp == 2 %}
     value.y = input[inputBase + 1u];
 {% endif %}
 {% endif %}
@@ -163,7 +162,7 @@ fn main(
   }
   storageBarrier();
-{% for stage in range(1, source.log2Length + 1) %}
   {
     const SPAN: u32 = 1u << {{ stage }}u;
     const HALF: u32 = SPAN >> 1u;
@@ -207,7 +206,7 @@ fn main(
   }
 {% endfor %}
-{% if source.inverse %}
   const SCALE: f32 = 1.0 / f32(N);
 {% elif USE_SCRATCH %}
   const SCALE: f32 = 1.0;
@@ -223,7 +222,7 @@ fn main(
     output[dst] = scratch[src] * SCALE;
     output[dst + 1u] = scratch[src + 1u] * SCALE;
   }
-{% elif source.inverse %}
   for (var k = tid; k < N; k = k + WG) {
     let index = transformBase + k * 2u;
     output[index] = output[index] * SCALE;

 // 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" %}
+{% set USE_SCRATCH = modeSpec != "dft" %}
 {% if not STFT %}
 // 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 %}
 {% if STFT %}
 const PI: f32 = 3.141592653589793;
 {% endif %}
+const N: u32 = {{ fftLength }}u;
 const N_HALF: u32 = N / 2u;
+const LOG2_N: u32 = {{ log2Length }}u;
 {% if USE_SCRATCH %}
+const NUM_BINS: u32 = {{ numBins }}u;
 {% endif %}
 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(
     var value = vec2<f32>(0.0);
 {% if STFT %}
     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 %}
 {% else %}
+    let inputBase = ((outerIndex * N + n) * params.inner + innerIndex) * {{ inComp }}u;
     value.x = input[inputBase];
+{% if inComp == 2 %}
     value.y = input[inputBase + 1u];
 {% endif %}
 {% endif %}
   }
   storageBarrier();
+{% for stage in range(1, log2Length + 1) %}
   {
     const SPAN: u32 = 1u << {{ stage }}u;
     const HALF: u32 = SPAN >> 1u;
   }
 {% endfor %}
+{% if inverse %}
   const SCALE: f32 = 1.0 / f32(N);
 {% elif USE_SCRATCH %}
   const SCALE: f32 = 1.0;
     output[dst] = scratch[src] * SCALE;
     output[dst + 1u] = scratch[src + 1u] * SCALE;
   }
+{% elif inverse %}
   for (var k = tid; k < N; k = k + WG) {
     let index = transformBase + k * 2u;
     output[index] = output[index] * SCALE;

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

@@ -1,17 +1,18 @@
 {{ 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,28 +85,20 @@ fn dft_precise_twiddle(phase: u32, period: u32, sign: f32) -> vec2<f32> {
 // 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;
-const NUM_BINS: u32 = {{ source.numBins }}u;
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 // 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" }};
 // Ping-pong scratch: stage reads [srcOff, srcOff+N), writes [dstOff, dstOff+N).
@@ -117,26 +110,23 @@ 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;
-{% if source.runtimeAxis is defined and source.runtimeAxis %}
-  // Dynamic-axis DFT pairs this axis-specialized pass with a general pass.
-  // The uniform guard ensures exactly one performs work for each invocation.
   let rawAxis = axis_in[0];
-  let normalizedAxis = select(rawAxis, rawAxis + {{ source.axisRank }}, rawAxis < 0);
-  if (normalizedAxis != {{ source.expectedAxis }}) { return; }
 {% endif %}
-  // 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 outerIndex = flat / params.inner;
   let innerIndex = flat - outerIndex * params.inner;
   for (var n = tid; n < N; n = n + WG) {
-    let base = ((outerIndex * N + n) * params.inner + innerIndex) * {{ source.inComp }}u;
     shRe[n] = input[base];
-{% if source.inComp == 2 %}
     shIm[n] = input[base + 1u];
 {% else %}
     shIm[n] = 0.0;
@@ -148,8 +138,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;
   for (var s = 0u; (1u << s) < N; s = s + 1u) {
@@ -178,9 +167,8 @@ 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.
-{% if source.inverse %}
   const SCALE: f32 = 1.0 / f32(N);
 {% else %}
   const SCALE: f32 = 1.0;

 {{ env.wgsl.resourceDeclarations }}
+{% set axisRank = axisRank | default(0) %}
+{% set expectedAxis = expectedAxis | default(0) %}
 // 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;
 // 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.
+// This specialization evaluates roots inline.
+// The input is [outer, N, inner, inComp] and the output is
 // [outer, numBins, inner, 2]; params.inner carries the (runtime) inner count
+// and dispatch.x = outer * inner workgroups. The inverse-plus-onesided
+// combination is outside this specialization.
+const N: u32 = {{ fftLength }}u;
 const N_HALF: u32 = N / 2u;
+const NUM_BINS: u32 = {{ numBins }}u;
 const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
 // Twiddle = exp(SIGN * i * 2*pi * k / span); SIGN flips for the inverse
+// transform.
+const SIGN: f32 = {{ "1.0" if inverse else "-1.0" }};
 // Ping-pong scratch: stage reads [srcOff, srcOff+N), writes [dstOff, dstOff+N).
         @builtin(local_invocation_id) lid: vec3<u32>,
         @builtin(num_workgroups) nwg: vec3<u32>) {
   let tid = lid.x;
+{% if runtimeAxis is defined and runtimeAxis %}
+  // This axis specialization returns uniformly unless the normalized dynamic
+  // axis equals its baked axis.
   let rawAxis = axis_in[0];
+  let normalizedAxis = select(rawAxis, rawAxis + {{ axisRank }}, rawAxis < 0);
+  if (normalizedAxis != {{ expectedAxis }}) { return; }
 {% endif %}
+  // 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 outerIndex = flat / params.inner;
   let innerIndex = flat - outerIndex * params.inner;
   for (var n = tid; n < N; n = n + WG) {
+    let base = ((outerIndex * N + n) * params.inner + innerIndex) * {{ inComp }}u;
     shRe[n] = input[base];
+{% if inComp == 2 %}
     shIm[n] = input[base + 1u];
 {% else %}
     shIm[n] = 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;
   for (var s = 0u; (1u << s) < N; s = s + 1u) {
   }
   // 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.
+{% if inverse %}
   const SCALE: f32 = 1.0 / f32(N);
 {% else %}
   const SCALE: f32 = 1.0;

build/webgpu/manifest.json CHANGED Viewed

@@ -2,521 +2,251 @@
   "domain": "ai.onnx",
   "name": "DFT",
   "sinceVersion": 20,
-  "description": "Computes the discrete Fourier Transform (DFT) of the input along a specified axis. For a signal of length N, output bin k is `sum_{n=0}^{N-1} exp(-2*pi*j*k*n/N) * x[n]`; the inverse divides by N and negates the exponent sign. Supports forward/inverse, real-to-complex (RFFT), and complex-to-real (IRFFT) modes via the `onesided` and `inverse` attributes.",
-  "inputs": [
-    {
-      "role": "input",
-      "dtype": "T",
-      "description": "Input signal tensor; the last dimension is 1 for real values or 2 for complex `(real, imaginary)` pairs."
-    },
-    {
-      "role": "dft_length",
-      "dtype": "L",
-      "optional": true,
-      "description": "Optional int32 scalar controlling the signal length used for the transform; input is zero-padded or truncated to this length.",
-      "rank": 0
-    },
-    {
-      "role": "axis",
-      "dtype": "I",
-      "optional": true,
-      "description": "Optional logical int64 scalar specifying the dimension over which to compute the DFT (the last axis is reserved for the real/imaginary component and is excluded); signed axes use int32 WebGPU storage and default to `-2` when omitted.",
-      "rank": 0
-    }
-  ],
-  "outputs": [
-    {
-      "role": "output",
-      "dtype": "T",
-      "description": "DFT result tensor; last dimension is 2 (complex) for forward DFT and RFFT, or 1 (real) for IRFFT.",
-      "rank": "ranks.input"
-    }
-  ],
-  "attributes": { "inverse": 0, "onesided": 0 },
-  "attributeDescriptions": {
-    "inverse": "When set to 1, computes the inverse DFT (IDFT/IRFFT) instead of the forward transform; default is 0 (forward).",
-    "onesided": "When set to 1, exploits conjugate symmetry to return only the non-redundant half of the spectrum (RFFT for forward, IRFFT for inverse); default is 0 (full spectrum)."
   },
   "attributeConstraints": { "inverse": { "values": [0, 1] }, "onesided": { "values": [0, 1] } },
   "typeConstraints": { "T": ["float32"], "L": ["int32"], "I": ["int64"] },
-  "args": {
-    "input": { "kind": "tensor", "semantic": "input", "role": "input" },
-    "dft_length": { "kind": "tensor", "semantic": "dft_length", "role": "dft_length", "required": false },
-    "axis": {
-      "kind": "tensor",
-      "semantic": "axis",
-      "role": "axis",
-      "dtype": "int32",
-      "narrowing": "checked",
-      "required": false
-    },
-    "output": { "kind": "tensor", "semantic": "output", "role": "output" }
   },
-  "tunables": { "WORKGROUP_SIZE": 128, "TILED_REAL_WORKGROUP_SIZE": 256, "TILED_REAL_MIN_LENGTH": 512 },
   "derive": {
     "dftAxisAttr": "ranks.input - 2",
     "dftAttrOuter": "outer(shapes.input, dftAxisAttr)",
     "dftAttrInputInner": "inner(shapes.input, dftAxisAttr) / dim(shapes.input, ranks.input - 1)",
     "dftAttrOutputOuter": "outer(shapes.output, dftAxisAttr)",
     "dftAttrOutputInner": "inner(shapes.output, dftAxisAttr) / dim(shapes.output, ranks.output - 1)",
-    "rank3BatchDispatchFits": "dim(shapes.input, 0) <= device.limits.maxComputeWorkgroupsPerDimension",
     "rank3Axis1SharedFftFits": "rank3BatchDispatchFits and pow2ceil(dim(shapes.input, 1)) == dim(shapes.input, 1) and dim(shapes.input, 1) >= 4 and dim(shapes.input, 1) * 16 <= device.limits.maxComputeWorkgroupStorageSize",
     "rank3Axis1StorageFftFits": "rank3BatchDispatchFits and pow2ceil(dim(shapes.input, 1)) == dim(shapes.input, 1) and dim(shapes.input, 1) >= 4 and dim(shapes.input, 1) * 16 > device.limits.maxComputeWorkgroupStorageSize",
     "rank3Axis1FftLog2": "log2ceil(dim(shapes.input, 1))",
-    "rank3Axis1ScratchFits": "dim(shapes.input, 0) * dim(shapes.input, 1) * 2 * 4 <= device.limits.maxStorageBufferBindingSize and dim(shapes.input, 0) * dim(shapes.input, 1) * 2 * 4 <= device.limits.maxBufferSize"
   },
-  "bindingSets": {
-    "contiguousTransform": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "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": "inner", "type": "u32", "value": 1 }] }
-      }
-    ],
-    "dftScratch": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "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": "inner", "type": "u32", "value": 1 }] }
-      }
-    ],
-    "fftLen": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "dft_len_in",
-        "arg": "dft_length",
-        "semantic": "dft_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32",
-        "length": 1
-      },
-      { "name": "output", "arg": "output", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "inner", "type": "u32", "value": 1 }] }
-      }
-    ],
-    "fftLenStatic": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "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": "inner", "type": "u32", "value": 1 }] }
-      }
-    ],
-    "fftAxis": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      {
-        "name": "axis_in",
-        "arg": "axis",
-        "semantic": "axis",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32",
-        "length": 1
-      },
-      { "name": "output", "arg": "output", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "inner", "type": "u32", "value": 1 }] }
-      }
-    ],
-    "sharedFftTransform": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "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": "inner",
-              "type": "u32",
-              "value": "inner(shapes.input, dftAxisAttr) / dim(shapes.input, ranks.input - 1)"
             }
-          ]
-        }
-      }
-    ],
-    "rank2Naive": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "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) / dim(shapes.output, 1)" },
-            { "name": "signalLength", "type": "u32", "value": "dim(shapes.input, 0)" },
-            { "name": "inputComponents", "type": "u32", "value": "dim(shapes.input, 1)" },
-            { "name": "outputComponents", "type": "u32", "value": "dim(shapes.output, 1)" },
-            { "name": "outputLength", "type": "u32", "value": "dim(shapes.output, 0)" },
-            { "name": "inverse", "type": "u32", "value": "attrs.inverse" },
-            { "name": "onesided", "type": "u32", "value": "attrs.onesided" }
-          ]
-        }
-      }
-    ],
-    "rank3Naive": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "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) / dim(shapes.output, 2)" },
-            { "name": "signalLength", "type": "u32", "value": "dim(shapes.input, 1)" },
-            { "name": "inputComponents", "type": "u32", "value": "dim(shapes.input, 2)" },
-            { "name": "outputComponents", "type": "u32", "value": "dim(shapes.output, 2)" },
-            { "name": "outputLength", "type": "u32", "value": "dim(shapes.output, 1)" },
-            { "name": "inverse", "type": "u32", "value": "attrs.inverse" },
-            { "name": "onesided", "type": "u32", "value": "attrs.onesided" }
-          ]
-        }
-      }
-    ],
-    "rank4Naive": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "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) / dim(shapes.output, 3)" },
-            { "name": "inputD1", "type": "u32", "value": "dim(shapes.input, 1)" },
-            { "name": "inputD2", "type": "u32", "value": "dim(shapes.input, 2)" },
-            { "name": "outputD1", "type": "u32", "value": "dim(shapes.output, 1)" },
-            { "name": "outputD2", "type": "u32", "value": "dim(shapes.output, 2)" },
-            { "name": "inputComponents", "type": "u32", "value": "dim(shapes.input, 3)" },
-            { "name": "outputComponents", "type": "u32", "value": "dim(shapes.output, 3)" },
-            { "name": "inverse", "type": "u32", "value": "attrs.inverse" },
-            { "name": "onesided", "type": "u32", "value": "attrs.onesided" }
-          ]
-        }
-      }
-    ],
-    "genNone": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "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) / dim(shapes.output, ranks.output - 1)" },
-            { "name": "rank", "type": "u32", "value": "4" },
-            { "name": "inputComponents", "type": "u32", "value": "dim(shapes.input, ranks.input - 1)" },
-            { "name": "outputComponents", "type": "u32", "value": "dim(shapes.output, ranks.output - 1)" },
-            { "name": "inverse", "type": "u32", "value": "attrs.inverse" },
-            { "name": "onesided", "type": "u32", "value": "attrs.onesided" },
-            { "name": "inD0", "type": "u32", "value": 1 },
-            { "name": "inD1", "type": "u32", "value": 1 },
-            { "name": "inD2", "type": "u32", "value": 1 },
-            { "name": "inD3", "type": "u32", "value": 1 },
-            { "name": "inD4", "type": "u32", "value": "dftAttrOuter" },
-            { "name": "inD5", "type": "u32", "value": "dim(shapes.input, dftAxisAttr)" },
-            { "name": "inD6", "type": "u32", "value": "dftAttrInputInner" },
-            { "name": "outD0", "type": "u32", "value": 1 },
-            { "name": "outD1", "type": "u32", "value": 1 },
-            { "name": "outD2", "type": "u32", "value": 1 },
-            { "name": "outD3", "type": "u32", "value": 1 },
-            { "name": "outD4", "type": "u32", "value": "dftAttrOutputOuter" },
-            { "name": "outD5", "type": "u32", "value": "dim(shapes.output, dftAxisAttr)" },
-            { "name": "outD6", "type": "u32", "value": "dftAttrOutputInner" },
-            { "name": "axisAttr", "type": "u32", "value": "1" },
-            { "name": "dftLengthAttr", "type": "u32", "value": 0 }
-          ]
-        }
-      }
-    ],
-    "genAxis": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "output", "arg": "output", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "axis_in",
-        "arg": "axis",
-        "semantic": "axis",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32",
-        "length": 1
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "count", "type": "u32", "value": "numel(shapes.output) / dim(shapes.output, ranks.output - 1)" },
-            { "name": "rank", "type": "u32", "value": "ranks.input" },
-            { "name": "inputComponents", "type": "u32", "value": "dim(shapes.input, ranks.input - 1)" },
-            { "name": "outputComponents", "type": "u32", "value": "dim(shapes.output, ranks.output - 1)" },
-            { "name": "inverse", "type": "u32", "value": "attrs.inverse" },
-            { "name": "onesided", "type": "u32", "value": "attrs.onesided" },
-            { "name": "inD0", "type": "u32", "value": "dim(shapes.input, ranks.input - 8) if ranks.input >= 8 else 1" },
-            { "name": "inD1", "type": "u32", "value": "dim(shapes.input, ranks.input - 7) if ranks.input >= 7 else 1" },
-            { "name": "inD2", "type": "u32", "value": "dim(shapes.input, ranks.input - 6) if ranks.input >= 6 else 1" },
-            { "name": "inD3", "type": "u32", "value": "dim(shapes.input, ranks.input - 5) if ranks.input >= 5 else 1" },
-            { "name": "inD4", "type": "u32", "value": "dim(shapes.input, ranks.input - 4) if ranks.input >= 4 else 1" },
-            { "name": "inD5", "type": "u32", "value": "dim(shapes.input, ranks.input - 3) if ranks.input >= 3 else 1" },
-            { "name": "inD6", "type": "u32", "value": "dim(shapes.input, ranks.input - 2)" },
-            {
-              "name": "outD0",
-              "type": "u32",
-              "value": "dim(shapes.output, ranks.output - 8) if ranks.output >= 8 else 1"
-            },
-            {
-              "name": "outD1",
-              "type": "u32",
-              "value": "dim(shapes.output, ranks.output - 7) if ranks.output >= 7 else 1"
-            },
-            {
-              "name": "outD2",
-              "type": "u32",
-              "value": "dim(shapes.output, ranks.output - 6) if ranks.output >= 6 else 1"
-            },
-            {
-              "name": "outD3",
-              "type": "u32",
-              "value": "dim(shapes.output, ranks.output - 5) if ranks.output >= 5 else 1"
-            },
-            {
-              "name": "outD4",
-              "type": "u32",
-              "value": "dim(shapes.output, ranks.output - 4) if ranks.output >= 4 else 1"
-            },
-            {
-              "name": "outD5",
-              "type": "u32",
-              "value": "dim(shapes.output, ranks.output - 3) if ranks.output >= 3 else 1"
-            },
-            { "name": "outD6", "type": "u32", "value": "dim(shapes.output, ranks.output - 2)" },
-            { "name": "dftLengthAttr", "type": "u32", "value": 0 }
-          ]
-        }
-      }
-    ],
-    "genLen": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "output", "arg": "output", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "dft_len_in",
-        "arg": "dft_length",
-        "semantic": "dft_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32",
-        "length": 1
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "count", "type": "u32", "value": "numel(shapes.output) / dim(shapes.output, ranks.output - 1)" },
-            { "name": "rank", "type": "u32", "value": "4" },
-            { "name": "inputComponents", "type": "u32", "value": "dim(shapes.input, ranks.input - 1)" },
-            { "name": "outputComponents", "type": "u32", "value": "dim(shapes.output, ranks.output - 1)" },
-            { "name": "inverse", "type": "u32", "value": "attrs.inverse" },
-            { "name": "onesided", "type": "u32", "value": "attrs.onesided" },
-            { "name": "inD0", "type": "u32", "value": 1 },
-            { "name": "inD1", "type": "u32", "value": 1 },
-            { "name": "inD2", "type": "u32", "value": 1 },
-            { "name": "inD3", "type": "u32", "value": 1 },
-            { "name": "inD4", "type": "u32", "value": "dftAttrOuter" },
-            { "name": "inD5", "type": "u32", "value": "dim(shapes.input, dftAxisAttr)" },
-            { "name": "inD6", "type": "u32", "value": "dftAttrInputInner" },
-            { "name": "outD0", "type": "u32", "value": 1 },
-            { "name": "outD1", "type": "u32", "value": 1 },
-            { "name": "outD2", "type": "u32", "value": 1 },
-            { "name": "outD3", "type": "u32", "value": 1 },
-            { "name": "outD4", "type": "u32", "value": "dftAttrOutputOuter" },
-            { "name": "outD5", "type": "u32", "value": "dim(shapes.output, dftAxisAttr)" },
-            { "name": "outD6", "type": "u32", "value": "dftAttrOutputInner" },
-            { "name": "axisAttr", "type": "u32", "value": "1" }
-          ]
-        }
-      }
-    ],
-    "genAxisLen": [
-      {
-        "name": "input",
-        "arg": "input",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "f32"
-      },
-      { "name": "output", "arg": "output", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "f32" },
-      {
-        "name": "dft_len_in",
-        "arg": "dft_length",
-        "semantic": "dft_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32",
-        "length": 1
-      },
-      {
-        "name": "axis_in",
-        "arg": "axis",
-        "semantic": "axis",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32",
-        "length": 1
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "count", "type": "u32", "value": "numel(shapes.output) / dim(shapes.output, ranks.output - 1)" },
-            { "name": "rank", "type": "u32", "value": "ranks.input" },
-            { "name": "inputComponents", "type": "u32", "value": "dim(shapes.input, ranks.input - 1)" },
-            { "name": "outputComponents", "type": "u32", "value": "dim(shapes.output, ranks.output - 1)" },
-            { "name": "inverse", "type": "u32", "value": "attrs.inverse" },
-            { "name": "onesided", "type": "u32", "value": "attrs.onesided" },
-            { "name": "inD0", "type": "u32", "value": "dim(shapes.input, ranks.input - 8) if ranks.input >= 8 else 1" },
-            { "name": "inD1", "type": "u32", "value": "dim(shapes.input, ranks.input - 7) if ranks.input >= 7 else 1" },
-            { "name": "inD2", "type": "u32", "value": "dim(shapes.input, ranks.input - 6) if ranks.input >= 6 else 1" },
-            { "name": "inD3", "type": "u32", "value": "dim(shapes.input, ranks.input - 5) if ranks.input >= 5 else 1" },
-            { "name": "inD4", "type": "u32", "value": "dim(shapes.input, ranks.input - 4) if ranks.input >= 4 else 1" },
-            { "name": "inD5", "type": "u32", "value": "dim(shapes.input, ranks.input - 3) if ranks.input >= 3 else 1" },
-            { "name": "inD6", "type": "u32", "value": "dim(shapes.input, ranks.input - 2)" },
-            {
-              "name": "outD0",
-              "type": "u32",
-              "value": "dim(shapes.output, ranks.output - 8) if ranks.output >= 8 else 1"
-            },
-            {
-              "name": "outD1",
-              "type": "u32",
-              "value": "dim(shapes.output, ranks.output - 7) if ranks.output >= 7 else 1"
-            },
-            {
-              "name": "outD2",
-              "type": "u32",
-              "value": "dim(shapes.output, ranks.output - 6) if ranks.output >= 6 else 1"
-            },
-            {
-              "name": "outD3",
-              "type": "u32",
-              "value": "dim(shapes.output, ranks.output - 5) if ranks.output >= 5 else 1"
-            },
-            {
-              "name": "outD4",
-              "type": "u32",
-              "value": "dim(shapes.output, ranks.output - 4) if ranks.output >= 4 else 1"
-            },
-            {
-              "name": "outD5",
-              "type": "u32",
-              "value": "dim(shapes.output, ranks.output - 3) if ranks.output >= 3 else 1"
-            },
-            { "name": "outD6", "type": "u32", "value": "dim(shapes.output, ranks.output - 2)" }
-          ]
         }
-      }
-    ]
-  },
-  "variants": [
     {
       "id": "batch_runtime_axis1_fft",
       "priority": 30,
@@ -525,33 +255,30 @@
         {
           "id": "axis1_fft",
           "name": "DFT.BatchRuntimeAxis1FftShared",
-          "source": {
-            "shader": "fft-stockham-shared.wgsl.jinja",
-            "inputs": {
-              "mode": "\"dft\"",
-              "fftLength": "dim(shapes.input, 1)",
-              "numBins": "dim(shapes.output, 1)",
-              "inComp": "dim(shapes.input, 2)",
-              "inverse": "attrs.inverse != 0",
-              "runtimeAxis": true,
-              "axisRank": 3,
-              "expectedAxis": 1
-            }
           },
-          "bindings": "fftAxis",
           "dispatch": { "x": "dim(shapes.input, 0)" }
         },
         {
           "id": "other_axis_general",
           "name": "DFT.BatchRuntimeOtherAxisGeneral",
-          "source": {
-            "shader": "dft-general.wgsl.jinja",
-            "inputs": { "hasAxisInput": true, "hasDftLengthInput": false, "skipRuntimeAxis1": true }
-          },
-          "bindings": "genAxis",
           "dispatch": {
-            "threads": "numel(shapes.output) / dim(shapes.output, ranks.output - 1)",
-            "workgroupSize": "tunables.WORKGROUP_SIZE"
           }
         }
       ]
@@ -567,18 +294,21 @@
         {
           "id": "main",
           "name": "DFT.BatchAxis1FftStorageOnesided",
-          "source": {
-            "shader": "fft-radix2-dit-storage.wgsl.jinja",
-            "inputs": {
-              "mode": "\"dft_scratch\"",
-              "fftLength": "dim(shapes.input, 1)",
-              "log2Length": "rank3Axis1FftLog2",
-              "numBins": "dim(shapes.output, 1)",
-              "inComp": 1,
-              "inverse": false
-            }
           },
-          "bindings": "dftScratch",
           "dispatch": { "x": "dim(shapes.input, 0)" }
         }
       ]
@@ -591,18 +321,16 @@
         {
           "id": "main",
           "name": "DFT.BatchAxis1FftStorageInplace",
-          "source": {
-            "shader": "fft-radix2-dit-storage.wgsl.jinja",
-            "inputs": {
-              "mode": "\"dft\"",
-              "fftLength": "dim(shapes.input, 1)",
-              "log2Length": "rank3Axis1FftLog2",
-              "numBins": "dim(shapes.output, 1)",
-              "inComp": "dim(shapes.input, 2)",
-              "inverse": "attrs.inverse != 0"
-            }
           },
-          "bindings": "contiguousTransform",
           "dispatch": { "x": "dim(shapes.input, 0)" }
         }
       ]
@@ -610,22 +338,20 @@
     {
       "id": "batch_axis1_tiled_real_nonpow2",
       "priority": 29,
       "requires": { "limits": { "maxComputeInvocationsPerWorkgroup": 256, "maxComputeWorkgroupSizeX": 256 } },
-      "when": ["not present.axis", "not present.dft_length", "ranks.input == 3", "ranks.output == 3", "dftAxisAttr == 1", "attrs.inverse == 0", "dim(shapes.input, 0) == dim(shapes.output, 0)", "dim(shapes.input, 0) > 0", "rank3BatchDispatchFits", "dim(shapes.input, 1) >= tunables.TILED_REAL_MIN_LENGTH", "pow2ceil(dim(shapes.input, 1)) != dim(shapes.input, 1)", "dim(shapes.input, 1) * 4 <= device.limits.maxComputeWorkgroupStorageSize", "dim(shapes.input, 2) == 1", "dim(shapes.output, 2) == 2", "((attrs.onesided == 0 and dim(shapes.output, 1) == dim(shapes.input, 1)) or (attrs.onesided != 0 and dim(shapes.output, 1) == floor(dim(shapes.input, 1) / 2) + 1))", "ceilDiv(dim(shapes.output, 1), tunables.TILED_REAL_WORKGROUP_SIZE) <= device.limits.maxComputeWorkgroupsPerDimension"],
       "passes": [
         {
           "id": "main",
           "name": "DFT.BatchAxis1TiledRealNonPowerOfTwo",
-          "source": {
-            "shader": "dft-tiled-real.wgsl.jinja",
-            "inputs": {
-              "mode": "\"dft\"",
-              "transformLength": "dim(shapes.input, 1)",
-              "numBins": "dim(shapes.output, 1)",
-              "workgroupSize": "tunables.TILED_REAL_WORKGROUP_SIZE"
-            }
           },
-          "bindings": "contiguousTransform",
           "dispatch": {
             "x": "ceilDiv(dim(shapes.output, 1), tunables.TILED_REAL_WORKGROUP_SIZE)",
             "y": "dim(shapes.input, 0)"
@@ -636,22 +362,33 @@
     {
       "id": "axis_fft_shared",
       "priority": 30,
-      "when": ["not present.axis", "not present.dft_length", "ranks.input == ranks.output", "((ranks.input == 2 and dftAxisAttr == 0) or (ranks.input == 3 and dftAxisAttr == 1 and dim(shapes.input, 0) == dim(shapes.output, 0)) or (ranks.input == 4 and dftAxisAttr == 2 and dim(shapes.input, 0) == dim(shapes.output, 0) and dim(shapes.input, 1) == dim(shapes.output, 1)))", "pow2ceil(dim(shapes.input, dftAxisAttr)) == dim(shapes.input, dftAxisAttr)", "dim(shapes.input, dftAxisAttr) >= 4", "dim(shapes.input, dftAxisAttr) * 16 <= device.limits.maxComputeWorkgroupStorageSize", "outer(shapes.input, dftAxisAttr) * inner(shapes.input, dftAxisAttr) / dim(shapes.input, ranks.input - 1) <= device.limits.maxComputeWorkgroupsPerDimension", "dim(shapes.input, ranks.input - 1) >= 1", "dim(shapes.input, ranks.input - 1) <= 2", "((attrs.onesided == 0 and dim(shapes.output, dftAxisAttr) == dim(shapes.input, dftAxisAttr) and dim(shapes.output, ranks.input - 1) == 2) or (attrs.onesided != 0 and attrs.inverse == 0 and dim(shapes.input, ranks.input - 1) == 1 and dim(shapes.output, dftAxisAttr) == floor(dim(shapes.input, dftAxisAttr) / 2) + 1 and dim(shapes.output, ranks.input - 1) == 2))"],
       "passes": [
         {
           "id": "main",
           "name": "DFT.AxisFftShared",
-          "source": {
-            "shader": "fft-stockham-shared.wgsl.jinja",
-            "inputs": {
-              "mode": "\"dft\"",
-              "fftLength": "dim(shapes.input, dftAxisAttr)",
-              "numBins": "dim(shapes.output, dftAxisAttr)",
-              "inComp": "dim(shapes.input, ranks.input - 1)",
-              "inverse": "attrs.inverse != 0"
-            }
           },
-          "bindings": "sharedFftTransform",
           "dispatch": {
             "x": "outer(shapes.input, dftAxisAttr) * inner(shapes.input, dftAxisAttr) / dim(shapes.input, ranks.input - 1)"
           }
@@ -666,17 +403,15 @@
         {
           "id": "main",
           "name": "DFT.BatchAxis1FftSharedExplicitLength",
-          "source": {
-            "shader": "fft-stockham-shared.wgsl.jinja",
-            "inputs": {
-              "mode": "\"dft\"",
-              "fftLength": "dim(shapes.input, 1)",
-              "numBins": "dim(shapes.output, 1)",
-              "inComp": "dim(shapes.input, 2)",
-              "inverse": "attrs.inverse != 0"
-            }
           },
-          "bindings": "fftLenStatic",
           "dispatch": { "x": "dim(shapes.input, 0)" }
         }
       ]
@@ -689,11 +424,28 @@
         {
           "id": "main",
           "name": "DFT.Rank2Axis0",
-          "source": { "shader": "dft-contiguous-naive.wgsl.jinja", "inputs": { "batched": false } },
-          "bindings": "rank2Naive",
           "dispatch": {
-            "threads": "numel(shapes.output) / dim(shapes.output, 1)",
-            "workgroupSize": "tunables.WORKGROUP_SIZE"
           }
         }
       ]
@@ -706,11 +458,28 @@
         {
           "id": "main",
           "name": "DFT",
-          "source": { "shader": "dft-contiguous-naive.wgsl.jinja", "inputs": { "batched": true } },
-          "bindings": "rank3Naive",
           "dispatch": {
-            "threads": "numel(shapes.output) / dim(shapes.output, 2)",
-            "workgroupSize": "tunables.WORKGROUP_SIZE"
           }
         }
       ]
@@ -724,10 +493,28 @@
           "id": "main",
           "name": "DFT.Rank4Axis2",
           "shader": "dft-rank4.wgsl.jinja",
-          "bindings": "rank4Naive",
           "dispatch": {
-            "threads": "numel(shapes.output) / dim(shapes.output, 3)",
-            "workgroupSize": "tunables.WORKGROUP_SIZE"
           }
         }
       ]
@@ -740,14 +527,13 @@
         {
           "id": "main",
           "name": "DFT.General",
-          "source": {
-            "shader": "dft-general.wgsl.jinja",
-            "inputs": { "hasAxisInput": false, "hasDftLengthInput": false, "skipRuntimeAxis1": false }
-          },
-          "bindings": "genNone",
           "dispatch": {
-            "threads": "numel(shapes.output) / dim(shapes.output, ranks.output - 1)",
-            "workgroupSize": "tunables.WORKGROUP_SIZE"
           }
         }
       ]
@@ -760,14 +546,13 @@
         {
           "id": "main",
           "name": "DFT.General",
-          "source": {
-            "shader": "dft-general.wgsl.jinja",
-            "inputs": { "hasAxisInput": true, "hasDftLengthInput": true, "skipRuntimeAxis1": false }
-          },
-          "bindings": "genAxisLen",
           "dispatch": {
-            "threads": "numel(shapes.output) / dim(shapes.output, ranks.output - 1)",
-            "workgroupSize": "tunables.WORKGROUP_SIZE"
           }
         }
       ]
@@ -780,14 +565,13 @@
         {
           "id": "main",
           "name": "DFT.General",
-          "source": {
-            "shader": "dft-general.wgsl.jinja",
-            "inputs": { "hasAxisInput": true, "hasDftLengthInput": false, "skipRuntimeAxis1": false }
-          },
-          "bindings": "genAxis",
           "dispatch": {
-            "threads": "numel(shapes.output) / dim(shapes.output, ranks.output - 1)",
-            "workgroupSize": "tunables.WORKGROUP_SIZE"
           }
         }
       ]
@@ -800,14 +584,13 @@
         {
           "id": "main",
           "name": "DFT.General",
-          "source": {
-            "shader": "dft-general.wgsl.jinja",
-            "inputs": { "hasAxisInput": false, "hasDftLengthInput": true, "skipRuntimeAxis1": false }
-          },
-          "bindings": "genLen",
           "dispatch": {
-            "threads": "numel(shapes.output) / dim(shapes.output, ranks.output - 1)",
-            "workgroupSize": "tunables.WORKGROUP_SIZE"
           }
         }
       ]

   "domain": "ai.onnx",
   "name": "DFT",
   "sinceVersion": 20,
+  "inputs": {
+    "input": { "dtype": "T" },
+    "dft_length": { "dtype": "L", "rank": 0, "optional": true },
+    "axis": { "dtype": "I", "rank": 0, "optional": true, "storage": "int32", "narrowing": "checked" }
   },
+  "outputs": { "output": { "dtype": "T", "rank": "ranks.input" } },
+  "attributes": { "inverse": { "default": 0 }, "onesided": { "default": 0 } },
   "attributeConstraints": { "inverse": { "values": [0, 1] }, "onesided": { "values": [0, 1] } },
   "typeConstraints": { "T": ["float32"], "L": ["int32"], "I": ["int64"] },
+  "tunables": {
+    "WORKGROUP_SIZE": { "default": 128 },
+    "TILED_REAL_WORKGROUP_SIZE": { "default": 256 },
+    "TILED_REAL_MIN_LENGTH": { "default": 512 },
+    "RUNTIME_AXIS_FFT_MIN_LENGTH": { "default": 64 }
   },
   "derive": {
     "dftAxisAttr": "ranks.input - 2",
     "dftAttrOuter": "outer(shapes.input, dftAxisAttr)",
     "dftAttrInputInner": "inner(shapes.input, dftAxisAttr) / dim(shapes.input, ranks.input - 1)",
     "dftAttrOutputOuter": "outer(shapes.output, dftAxisAttr)",
     "dftAttrOutputInner": "inner(shapes.output, dftAxisAttr) / dim(shapes.output, ranks.output - 1)",
+    "rank3BatchDispatchFits": "dim(shapes.input, 0) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "rank3Axis1SharedFftFits": "rank3BatchDispatchFits and pow2ceil(dim(shapes.input, 1)) == dim(shapes.input, 1) and dim(shapes.input, 1) >= 4 and dim(shapes.input, 1) * 16 <= device.limits.maxComputeWorkgroupStorageSize",
     "rank3Axis1StorageFftFits": "rank3BatchDispatchFits and pow2ceil(dim(shapes.input, 1)) == dim(shapes.input, 1) and dim(shapes.input, 1) >= 4 and dim(shapes.input, 1) * 16 > device.limits.maxComputeWorkgroupStorageSize",
     "rank3Axis1FftLog2": "log2ceil(dim(shapes.input, 1))",
+    "rank3Axis1ScratchFits": "dim(shapes.input, 0) * dim(shapes.input, 1) * 2 * 4 <= device.limits.maxStorageBufferBindingSize and dim(shapes.input, 0) * dim(shapes.input, 1) * 2 * 4 <= device.limits.maxBufferSize",
+    "fftAnyAxisLen0": "dim(shapes.input, 0) if ranks.input - 1 > 0 else 0",
+    "fftAnyAxisLen1": "dim(shapes.input, 1) if ranks.input - 1 > 1 else 0",
+    "fftAnyAxisLen2": "dim(shapes.input, 2) if ranks.input - 1 > 2 else 0",
+    "fftAnyAxisLen3": "dim(shapes.input, 3) if ranks.input - 1 > 3 else 0",
+    "fftAnyAxisLen4": "dim(shapes.input, 4) if ranks.input - 1 > 4 else 0",
+    "fftAnyAxisLen5": "dim(shapes.input, 5) if ranks.input - 1 > 5 else 0",
+    "fftAnyAxisLen6": "dim(shapes.input, 6) if ranks.input - 1 > 6 else 0",
+    "fftAnyAxisCand0": "fftAnyAxisLen0 if (fftAnyAxisLen0 >= tunables.RUNTIME_AXIS_FFT_MIN_LENGTH and pow2ceil(fftAnyAxisLen0) == fftAnyAxisLen0 and fftAnyAxisLen0 * 16 <= device.limits.maxComputeWorkgroupStorageSize) else 0",
+    "fftAnyAxisCand1": "fftAnyAxisLen1 if (fftAnyAxisLen1 >= tunables.RUNTIME_AXIS_FFT_MIN_LENGTH and pow2ceil(fftAnyAxisLen1) == fftAnyAxisLen1 and fftAnyAxisLen1 * 16 <= device.limits.maxComputeWorkgroupStorageSize) else 0",
+    "fftAnyAxisCand2": "fftAnyAxisLen2 if (fftAnyAxisLen2 >= tunables.RUNTIME_AXIS_FFT_MIN_LENGTH and pow2ceil(fftAnyAxisLen2) == fftAnyAxisLen2 and fftAnyAxisLen2 * 16 <= device.limits.maxComputeWorkgroupStorageSize) else 0",
+    "fftAnyAxisCand3": "fftAnyAxisLen3 if (fftAnyAxisLen3 >= tunables.RUNTIME_AXIS_FFT_MIN_LENGTH and pow2ceil(fftAnyAxisLen3) == fftAnyAxisLen3 and fftAnyAxisLen3 * 16 <= device.limits.maxComputeWorkgroupStorageSize) else 0",
+    "fftAnyAxisCand4": "fftAnyAxisLen4 if (fftAnyAxisLen4 >= tunables.RUNTIME_AXIS_FFT_MIN_LENGTH and pow2ceil(fftAnyAxisLen4) == fftAnyAxisLen4 and fftAnyAxisLen4 * 16 <= device.limits.maxComputeWorkgroupStorageSize) else 0",
+    "fftAnyAxisCand5": "fftAnyAxisLen5 if (fftAnyAxisLen5 >= tunables.RUNTIME_AXIS_FFT_MIN_LENGTH and pow2ceil(fftAnyAxisLen5) == fftAnyAxisLen5 and fftAnyAxisLen5 * 16 <= device.limits.maxComputeWorkgroupStorageSize) else 0",
+    "fftAnyAxisCand6": "fftAnyAxisLen6 if (fftAnyAxisLen6 >= tunables.RUNTIME_AXIS_FFT_MIN_LENGTH and pow2ceil(fftAnyAxisLen6) == fftAnyAxisLen6 and fftAnyAxisLen6 * 16 <= device.limits.maxComputeWorkgroupStorageSize) else 0",
+    "fftAnyAxisMaxLen": "max(fftAnyAxisCand0, max(fftAnyAxisCand1, max(fftAnyAxisCand2, max(fftAnyAxisCand3, max(fftAnyAxisCand4, max(fftAnyAxisCand5, fftAnyAxisCand6))))))",
+    "fftAnyAxisMinLen": "min((fftAnyAxisCand0 if fftAnyAxisCand0 > 0 else 1073741824), min((fftAnyAxisCand1 if fftAnyAxisCand1 > 0 else 1073741824), min((fftAnyAxisCand2 if fftAnyAxisCand2 > 0 else 1073741824), min((fftAnyAxisCand3 if fftAnyAxisCand3 > 0 else 1073741824), min((fftAnyAxisCand4 if fftAnyAxisCand4 > 0 else 1073741824), min((fftAnyAxisCand5 if fftAnyAxisCand5 > 0 else 1073741824), (fftAnyAxisCand6 if fftAnyAxisCand6 > 0 else 1073741824)))))))",
+    "fftAnyAxisLogicalCount": "numel(shapes.input) / dim(shapes.input, ranks.input - 1) if ranks.input >= 1 and dim(shapes.input, ranks.input - 1) > 0 else 0",
+    "fftAnyAxisTransforms": "fftAnyAxisLogicalCount / fftAnyAxisMinLen if fftAnyAxisMaxLen >= tunables.RUNTIME_AXIS_FFT_MIN_LENGTH else 0"
+  },
+  "bindings": {
+    "input": { "buffer": "read-only-storage", "elementType": "f32" },
+    "axis_in": { "arg": "axis", "buffer": "read-only-storage", "elementType": "i32", "length": 1 },
+    "output": { "buffer": "storage", "elementType": "f32" },
+    "params_2": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "count", "type": "u32", "value": "numel(shapes.output) / dim(shapes.output, ranks.output - 1)" },
+        { "name": "rank", "type": "u32", "value": "ranks.input" },
+        { "name": "inputComponents", "type": "u32", "value": "dim(shapes.input, ranks.input - 1)" },
+        { "name": "outputComponents", "type": "u32", "value": "dim(shapes.output, ranks.output - 1)" },
+        { "name": "inverse", "type": "u32", "value": "attrs.inverse" },
+        { "name": "onesided", "type": "u32", "value": "attrs.onesided" },
+        { "name": "inD0", "type": "u32", "value": "dim(shapes.input, ranks.input - 8) if ranks.input >= 8 else 1" },
+        { "name": "inD1", "type": "u32", "value": "dim(shapes.input, ranks.input - 7) if ranks.input >= 7 else 1" },
+        { "name": "inD2", "type": "u32", "value": "dim(shapes.input, ranks.input - 6) if ranks.input >= 6 else 1" },
+        { "name": "inD3", "type": "u32", "value": "dim(shapes.input, ranks.input - 5) if ranks.input >= 5 else 1" },
+        { "name": "inD4", "type": "u32", "value": "dim(shapes.input, ranks.input - 4) if ranks.input >= 4 else 1" },
+        { "name": "inD5", "type": "u32", "value": "dim(shapes.input, ranks.input - 3) if ranks.input >= 3 else 1" },
+        { "name": "inD6", "type": "u32", "value": "dim(shapes.input, ranks.input - 2)" },
+        { "name": "outD0", "type": "u32", "value": "dim(shapes.output, ranks.output - 8) if ranks.output >= 8 else 1" },
+        { "name": "outD1", "type": "u32", "value": "dim(shapes.output, ranks.output - 7) if ranks.output >= 7 else 1" },
+        { "name": "outD2", "type": "u32", "value": "dim(shapes.output, ranks.output - 6) if ranks.output >= 6 else 1" },
+        { "name": "outD3", "type": "u32", "value": "dim(shapes.output, ranks.output - 5) if ranks.output >= 5 else 1" },
+        { "name": "outD4", "type": "u32", "value": "dim(shapes.output, ranks.output - 4) if ranks.output >= 4 else 1" },
+        { "name": "outD5", "type": "u32", "value": "dim(shapes.output, ranks.output - 3) if ranks.output >= 3 else 1" },
+        { "name": "outD6", "type": "u32", "value": "dim(shapes.output, ranks.output - 2)" },
+        { "name": "dftLengthAttr", "type": "u32", "value": 0 }
+      ]
+    },
+    "params": { "buffer": "uniform", "struct": [{ "name": "inner", "type": "u32", "value": 1 }] },
+    "params_8": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "count", "type": "u32", "value": "numel(shapes.output) / dim(shapes.output, ranks.output - 1)" },
+        { "name": "rank", "type": "u32", "value": "4" },
+        { "name": "inputComponents", "type": "u32", "value": "dim(shapes.input, ranks.input - 1)" },
+        { "name": "outputComponents", "type": "u32", "value": "dim(shapes.output, ranks.output - 1)" },
+        { "name": "inverse", "type": "u32", "value": "attrs.inverse" },
+        { "name": "onesided", "type": "u32", "value": "attrs.onesided" },
+        { "name": "inD0", "type": "u32", "value": 1 },
+        { "name": "inD1", "type": "u32", "value": 1 },
+        { "name": "inD2", "type": "u32", "value": 1 },
+        { "name": "inD3", "type": "u32", "value": 1 },
+        { "name": "inD4", "type": "u32", "value": "dftAttrOuter" },
+        { "name": "inD5", "type": "u32", "value": "dim(shapes.input, dftAxisAttr)" },
+        { "name": "inD6", "type": "u32", "value": "dftAttrInputInner" },
+        { "name": "outD0", "type": "u32", "value": 1 },
+        { "name": "outD1", "type": "u32", "value": 1 },
+        { "name": "outD2", "type": "u32", "value": 1 },
+        { "name": "outD3", "type": "u32", "value": 1 },
+        { "name": "outD4", "type": "u32", "value": "dftAttrOutputOuter" },
+        { "name": "outD5", "type": "u32", "value": "dim(shapes.output, dftAxisAttr)" },
+        { "name": "outD6", "type": "u32", "value": "dftAttrOutputInner" },
+        { "name": "axisAttr", "type": "u32", "value": "1" },
+        { "name": "dftLengthAttr", "type": "u32", "value": 0 }
+      ]
+    },
+    "dft_len_in": { "arg": "dft_length", "buffer": "read-only-storage", "elementType": "i32", "length": 1 },
+    "params_9": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "count", "type": "u32", "value": "numel(shapes.output) / dim(shapes.output, ranks.output - 1)" },
+        { "name": "rank", "type": "u32", "value": "ranks.input" },
+        { "name": "inputComponents", "type": "u32", "value": "dim(shapes.input, ranks.input - 1)" },
+        { "name": "outputComponents", "type": "u32", "value": "dim(shapes.output, ranks.output - 1)" },
+        { "name": "inverse", "type": "u32", "value": "attrs.inverse" },
+        { "name": "onesided", "type": "u32", "value": "attrs.onesided" },
+        { "name": "inD0", "type": "u32", "value": "dim(shapes.input, ranks.input - 8) if ranks.input >= 8 else 1" },
+        { "name": "inD1", "type": "u32", "value": "dim(shapes.input, ranks.input - 7) if ranks.input >= 7 else 1" },
+        { "name": "inD2", "type": "u32", "value": "dim(shapes.input, ranks.input - 6) if ranks.input >= 6 else 1" },
+        { "name": "inD3", "type": "u32", "value": "dim(shapes.input, ranks.input - 5) if ranks.input >= 5 else 1" },
+        { "name": "inD4", "type": "u32", "value": "dim(shapes.input, ranks.input - 4) if ranks.input >= 4 else 1" },
+        { "name": "inD5", "type": "u32", "value": "dim(shapes.input, ranks.input - 3) if ranks.input >= 3 else 1" },
+        { "name": "inD6", "type": "u32", "value": "dim(shapes.input, ranks.input - 2)" },
+        { "name": "outD0", "type": "u32", "value": "dim(shapes.output, ranks.output - 8) if ranks.output >= 8 else 1" },
+        { "name": "outD1", "type": "u32", "value": "dim(shapes.output, ranks.output - 7) if ranks.output >= 7 else 1" },
+        { "name": "outD2", "type": "u32", "value": "dim(shapes.output, ranks.output - 6) if ranks.output >= 6 else 1" },
+        { "name": "outD3", "type": "u32", "value": "dim(shapes.output, ranks.output - 5) if ranks.output >= 5 else 1" },
+        { "name": "outD4", "type": "u32", "value": "dim(shapes.output, ranks.output - 4) if ranks.output >= 4 else 1" },
+        { "name": "outD5", "type": "u32", "value": "dim(shapes.output, ranks.output - 3) if ranks.output >= 3 else 1" },
+        { "name": "outD6", "type": "u32", "value": "dim(shapes.output, ranks.output - 2)" }
+      ]
+    },
+    "params_10": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "count", "type": "u32", "value": "numel(shapes.output) / dim(shapes.output, ranks.output - 1)" },
+        { "name": "rank", "type": "u32", "value": "4" },
+        { "name": "inputComponents", "type": "u32", "value": "dim(shapes.input, ranks.input - 1)" },
+        { "name": "outputComponents", "type": "u32", "value": "dim(shapes.output, ranks.output - 1)" },
+        { "name": "inverse", "type": "u32", "value": "attrs.inverse" },
+        { "name": "onesided", "type": "u32", "value": "attrs.onesided" },
+        { "name": "inD0", "type": "u32", "value": 1 },
+        { "name": "inD1", "type": "u32", "value": 1 },
+        { "name": "inD2", "type": "u32", "value": 1 },
+        { "name": "inD3", "type": "u32", "value": 1 },
+        { "name": "inD4", "type": "u32", "value": "dftAttrOuter" },
+        { "name": "inD5", "type": "u32", "value": "dim(shapes.input, dftAxisAttr)" },
+        { "name": "inD6", "type": "u32", "value": "dftAttrInputInner" },
+        { "name": "outD0", "type": "u32", "value": 1 },
+        { "name": "outD1", "type": "u32", "value": 1 },
+        { "name": "outD2", "type": "u32", "value": 1 },
+        { "name": "outD3", "type": "u32", "value": 1 },
+        { "name": "outD4", "type": "u32", "value": "dftAttrOutputOuter" },
+        { "name": "outD5", "type": "u32", "value": "dim(shapes.output, dftAxisAttr)" },
+        { "name": "outD6", "type": "u32", "value": "dftAttrOutputInner" },
+        { "name": "axisAttr", "type": "u32", "value": "1" }
+      ]
+    }
   },
+  "variants": [
+    {
+      "id": "runtime_axis_shared_fft",
+      "priority": 6,
+      "when": ["present.axis", "numel(shapes.axis) == 1", "not present.dft_length", "attrs.onesided == 0", "ranks.input >= 2", "ranks.input <= 8", "ranks.input == ranks.output", "ranks.input != 3", "dim(shapes.input, ranks.input - 1) >= 1", "dim(shapes.input, ranks.input - 1) <= 2", "dim(shapes.output, ranks.output - 1) == 2", "numel(shapes.output) == 2 * fftAnyAxisLogicalCount", "fftAnyAxisMaxLen >= tunables.RUNTIME_AXIS_FFT_MIN_LENGTH", "tunables.RUNTIME_AXIS_FFT_MIN_LENGTH >= 4", "pow2ceil(tunables.RUNTIME_AXIS_FFT_MIN_LENGTH) == tunables.RUNTIME_AXIS_FFT_MIN_LENGTH", "fftAnyAxisMaxLen * 16 <= device.limits.maxComputeWorkgroupStorageSize", "ceilDiv(fftAnyAxisTransforms, min(device.limits.maxComputeWorkgroupsPerDimension, 65535)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
+      "passes": [
+        {
+          "id": "shared_fft",
+          "name": "DFT.RuntimeAxisFftShared",
+          "shader": "dft-runtime-axis-fft-shared.wgsl.jinja",
+          "derive": {
+            "fftMaxLength": "fftAnyAxisMaxLen",
+            "fftMinLength": "tunables.RUNTIME_AXIS_FFT_MIN_LENGTH",
+            "inComp": "dim(shapes.input, ranks.input - 1)",
+            "inverse": "attrs.inverse != 0"
+          },
+          "bindings": [
+            "input",
+            "axis_in",
+            "output",
             {
+              "name": "params",
+              "struct": [
+                { "name": "rank", "type": "u32", "value": "ranks.input" },
+                {
+                  "name": "inD0",
+                  "type": "u32",
+                  "value": "dim(shapes.input, ranks.input - 8) if ranks.input >= 8 else 1"
+                },
+                {
+                  "name": "inD1",
+                  "type": "u32",
+                  "value": "dim(shapes.input, ranks.input - 7) if ranks.input >= 7 else 1"
+                },
+                {
+                  "name": "inD2",
+                  "type": "u32",
+                  "value": "dim(shapes.input, ranks.input - 6) if ranks.input >= 6 else 1"
+                },
+                {
+                  "name": "inD3",
+                  "type": "u32",
+                  "value": "dim(shapes.input, ranks.input - 5) if ranks.input >= 5 else 1"
+                },
+                {
+                  "name": "inD4",
+                  "type": "u32",
+                  "value": "dim(shapes.input, ranks.input - 4) if ranks.input >= 4 else 1"
+                },
+                {
+                  "name": "inD5",
+                  "type": "u32",
+                  "value": "dim(shapes.input, ranks.input - 3) if ranks.input >= 3 else 1"
+                },
+                {
+                  "name": "inD6",
+                  "type": "u32",
+                  "value": "dim(shapes.input, ranks.input - 2) if ranks.input >= 2 else 1"
+                }
+              ]
             }
+          ],
+          "dispatch": { "x": "min(fftAnyAxisTransforms, 65535)", "y": "ceilDiv(fftAnyAxisTransforms, 65535)", "z": 1 }
+        },
+        {
+          "id": "other_axis_general",
+          "name": "DFT.RuntimeAxisOtherGeneral",
+          "shader": "dft-general.wgsl.jinja",
+          "derive": {
+            "hasAxisInput": true,
+            "hasDftLengthInput": false,
+            "skipRuntimeAxis1": false,
+            "skipSharedFftAxes": true,
+            "sharedFftMinLength": "tunables.RUNTIME_AXIS_FFT_MIN_LENGTH",
+            "sharedFftMaxLength": "fftAnyAxisMaxLen"
+          },
+          "bindings": ["input", "output", "axis_in", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.output) / dim(shapes.output, ranks.output - 1)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.output) / dim(shapes.output, ranks.output - 1)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
+      ]
+    },
     {
       "id": "batch_runtime_axis1_fft",
       "priority": 30,
         {
           "id": "axis1_fft",
           "name": "DFT.BatchRuntimeAxis1FftShared",
+          "shader": "fft-stockham-shared.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"dft\"",
+            "fftLength": "dim(shapes.input, 1)",
+            "numBins": "dim(shapes.output, 1)",
+            "inComp": "dim(shapes.input, 2)",
+            "inverse": "attrs.inverse != 0",
+            "runtimeAxis": true,
+            "axisRank": 3,
+            "expectedAxis": 1
           },
+          "bindings": ["input", "axis_in", "output", "params"],
           "dispatch": { "x": "dim(shapes.input, 0)" }
         },
         {
           "id": "other_axis_general",
           "name": "DFT.BatchRuntimeOtherAxisGeneral",
+          "shader": "dft-general.wgsl.jinja",
+          "derive": { "hasAxisInput": true, "hasDftLengthInput": false, "skipRuntimeAxis1": true },
+          "bindings": ["input", "output", "axis_in", "params_2"],
           "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.output) / dim(shapes.output, ranks.output - 1)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.output) / dim(shapes.output, ranks.output - 1)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         }
       ]
         {
           "id": "main",
           "name": "DFT.BatchAxis1FftStorageOnesided",
+          "shader": "fft-radix2-dit-storage.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"dft_scratch\"",
+            "fftLength": "dim(shapes.input, 1)",
+            "log2Length": "rank3Axis1FftLog2",
+            "numBins": "dim(shapes.output, 1)",
+            "inComp": 1,
+            "inverse": false
           },
+          "bindings": [
+            "input",
+            { "scratch": "fft_scratch", "name": "scratch", "elementType": "f32" },
+            "output",
+            "params"
+          ],
           "dispatch": { "x": "dim(shapes.input, 0)" }
         }
       ]
         {
           "id": "main",
           "name": "DFT.BatchAxis1FftStorageInplace",
+          "shader": "fft-radix2-dit-storage.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"dft\"",
+            "fftLength": "dim(shapes.input, 1)",
+            "log2Length": "rank3Axis1FftLog2",
+            "numBins": "dim(shapes.output, 1)",
+            "inComp": "dim(shapes.input, 2)",
+            "inverse": "attrs.inverse != 0"
           },
+          "bindings": ["input", "output", "params"],
           "dispatch": { "x": "dim(shapes.input, 0)" }
         }
       ]
     {
       "id": "batch_axis1_tiled_real_nonpow2",
       "priority": 29,
+      "when": ["not present.axis", "not present.dft_length", "ranks.input == 3", "ranks.output == 3", "dftAxisAttr == 1", "attrs.inverse == 0", "dim(shapes.input, 0) == dim(shapes.output, 0)", "dim(shapes.input, 0) > 0", "rank3BatchDispatchFits", "dim(shapes.input, 1) >= tunables.TILED_REAL_MIN_LENGTH", "pow2ceil(dim(shapes.input, 1)) != dim(shapes.input, 1)", "dim(shapes.input, 1) * 4 <= device.limits.maxComputeWorkgroupStorageSize", "dim(shapes.input, 2) == 1", "dim(shapes.output, 2) == 2", "((attrs.onesided == 0 and dim(shapes.output, 1) == dim(shapes.input, 1)) or (attrs.onesided != 0 and dim(shapes.output, 1) == floor(dim(shapes.input, 1) / 2) + 1))", "ceilDiv(dim(shapes.output, 1), tunables.TILED_REAL_WORKGROUP_SIZE) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"],
       "requires": { "limits": { "maxComputeInvocationsPerWorkgroup": 256, "maxComputeWorkgroupSizeX": 256 } },
       "passes": [
         {
           "id": "main",
           "name": "DFT.BatchAxis1TiledRealNonPowerOfTwo",
+          "shader": "dft-tiled-real.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"dft\"",
+            "transformLength": "dim(shapes.input, 1)",
+            "numBins": "dim(shapes.output, 1)",
+            "workgroupSizeSpec": "tunables.TILED_REAL_WORKGROUP_SIZE"
           },
+          "bindings": ["input", "output", "params"],
           "dispatch": {
             "x": "ceilDiv(dim(shapes.output, 1), tunables.TILED_REAL_WORKGROUP_SIZE)",
             "y": "dim(shapes.input, 0)"
     {
       "id": "axis_fft_shared",
       "priority": 30,
+      "when": ["not present.axis", "not present.dft_length", "ranks.input == ranks.output", "((ranks.input == 2 and dftAxisAttr == 0) or (ranks.input == 3 and dftAxisAttr == 1 and dim(shapes.input, 0) == dim(shapes.output, 0)) or (ranks.input == 4 and dftAxisAttr == 2 and dim(shapes.input, 0) == dim(shapes.output, 0) and dim(shapes.input, 1) == dim(shapes.output, 1)))", "pow2ceil(dim(shapes.input, dftAxisAttr)) == dim(shapes.input, dftAxisAttr)", "dim(shapes.input, dftAxisAttr) >= 4", "dim(shapes.input, dftAxisAttr) * 16 <= device.limits.maxComputeWorkgroupStorageSize", "outer(shapes.input, dftAxisAttr) * inner(shapes.input, dftAxisAttr) / dim(shapes.input, ranks.input - 1) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "dim(shapes.input, ranks.input - 1) >= 1", "dim(shapes.input, ranks.input - 1) <= 2", "((attrs.onesided == 0 and dim(shapes.output, dftAxisAttr) == dim(shapes.input, dftAxisAttr) and dim(shapes.output, ranks.input - 1) == 2) or (attrs.onesided != 0 and attrs.inverse == 0 and dim(shapes.input, ranks.input - 1) == 1 and dim(shapes.output, dftAxisAttr) == floor(dim(shapes.input, dftAxisAttr) / 2) + 1 and dim(shapes.output, ranks.input - 1) == 2))"],
       "passes": [
         {
           "id": "main",
           "name": "DFT.AxisFftShared",
+          "shader": "fft-stockham-shared.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"dft\"",
+            "fftLength": "dim(shapes.input, dftAxisAttr)",
+            "numBins": "dim(shapes.output, dftAxisAttr)",
+            "inComp": "dim(shapes.input, ranks.input - 1)",
+            "inverse": "attrs.inverse != 0"
           },
+          "bindings": [
+            "input",
+            "output",
+            {
+              "name": "params",
+              "struct": [
+                {
+                  "name": "inner",
+                  "type": "u32",
+                  "value": "inner(shapes.input, dftAxisAttr) / dim(shapes.input, ranks.input - 1)"
+                }
+              ]
+            }
+          ],
           "dispatch": {
             "x": "outer(shapes.input, dftAxisAttr) * inner(shapes.input, dftAxisAttr) / dim(shapes.input, ranks.input - 1)"
           }
         {
           "id": "main",
           "name": "DFT.BatchAxis1FftSharedExplicitLength",
+          "shader": "fft-stockham-shared.wgsl.jinja",
+          "derive": {
+            "modeSpec": "\"dft\"",
+            "fftLength": "dim(shapes.input, 1)",
+            "numBins": "dim(shapes.output, 1)",
+            "inComp": "dim(shapes.input, 2)",
+            "inverse": "attrs.inverse != 0"
           },
+          "bindings": ["input", "output", "params"],
           "dispatch": { "x": "dim(shapes.input, 0)" }
         }
       ]
         {
           "id": "main",
           "name": "DFT.Rank2Axis0",
+          "shader": "dft-contiguous-naive.wgsl.jinja",
+          "derive": { "batched": false },
+          "bindings": [
+            "input",
+            "output",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "count", "type": "u32", "value": "numel(shapes.output) / dim(shapes.output, 1)" },
+                { "name": "signalLength", "type": "u32", "value": "dim(shapes.input, 0)" },
+                { "name": "inputComponents", "type": "u32", "value": "dim(shapes.input, 1)" },
+                { "name": "outputComponents", "type": "u32", "value": "dim(shapes.output, 1)" },
+                { "name": "outputLength", "type": "u32", "value": "dim(shapes.output, 0)" },
+                { "name": "inverse", "type": "u32", "value": "attrs.inverse" },
+                { "name": "onesided", "type": "u32", "value": "attrs.onesided" }
+              ]
+            }
+          ],
           "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.output) / dim(shapes.output, 1)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.output) / dim(shapes.output, 1)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         }
       ]
         {
           "id": "main",
           "name": "DFT",
+          "shader": "dft-contiguous-naive.wgsl.jinja",
+          "derive": { "batched": true },
+          "bindings": [
+            "input",
+            "output",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "count", "type": "u32", "value": "numel(shapes.output) / dim(shapes.output, 2)" },
+                { "name": "signalLength", "type": "u32", "value": "dim(shapes.input, 1)" },
+                { "name": "inputComponents", "type": "u32", "value": "dim(shapes.input, 2)" },
+                { "name": "outputComponents", "type": "u32", "value": "dim(shapes.output, 2)" },
+                { "name": "outputLength", "type": "u32", "value": "dim(shapes.output, 1)" },
+                { "name": "inverse", "type": "u32", "value": "attrs.inverse" },
+                { "name": "onesided", "type": "u32", "value": "attrs.onesided" }
+              ]
+            }
+          ],
           "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.output) / dim(shapes.output, 2)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.output) / dim(shapes.output, 2)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         }
       ]
           "id": "main",
           "name": "DFT.Rank4Axis2",
           "shader": "dft-rank4.wgsl.jinja",
+          "bindings": [
+            "input",
+            "output",
+            {
+              "name": "params",
+              "struct": [
+                { "name": "count", "type": "u32", "value": "numel(shapes.output) / dim(shapes.output, 3)" },
+                { "name": "inputD1", "type": "u32", "value": "dim(shapes.input, 1)" },
+                { "name": "inputD2", "type": "u32", "value": "dim(shapes.input, 2)" },
+                { "name": "outputD1", "type": "u32", "value": "dim(shapes.output, 1)" },
+                { "name": "outputD2", "type": "u32", "value": "dim(shapes.output, 2)" },
+                { "name": "inputComponents", "type": "u32", "value": "dim(shapes.input, 3)" },
+                { "name": "outputComponents", "type": "u32", "value": "dim(shapes.output, 3)" },
+                { "name": "inverse", "type": "u32", "value": "attrs.inverse" },
+                { "name": "onesided", "type": "u32", "value": "attrs.onesided" }
+              ]
+            }
+          ],
           "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.output) / dim(shapes.output, 3)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.output) / dim(shapes.output, 3)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         }
       ]
         {
           "id": "main",
           "name": "DFT.General",
+          "shader": "dft-general.wgsl.jinja",
+          "derive": { "hasAxisInput": false, "hasDftLengthInput": false, "skipRuntimeAxis1": false },
+          "bindings": ["input", "output", "params_8"],
           "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.output) / dim(shapes.output, ranks.output - 1)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.output) / dim(shapes.output, ranks.output - 1)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         }
       ]
         {
           "id": "main",
           "name": "DFT.General",
+          "shader": "dft-general.wgsl.jinja",
+          "derive": { "hasAxisInput": true, "hasDftLengthInput": true, "skipRuntimeAxis1": false },
+          "bindings": ["input", "output", "dft_len_in", "axis_in", "params_9"],
           "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.output) / dim(shapes.output, ranks.output - 1)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.output) / dim(shapes.output, ranks.output - 1)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         }
       ]
         {
           "id": "main",
           "name": "DFT.General",
+          "shader": "dft-general.wgsl.jinja",
+          "derive": { "hasAxisInput": true, "hasDftLengthInput": false, "skipRuntimeAxis1": false },
+          "bindings": ["input", "output", "axis_in", "params_2"],
           "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.output) / dim(shapes.output, ranks.output - 1)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.output) / dim(shapes.output, ranks.output - 1)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         }
       ]
         {
           "id": "main",
           "name": "DFT.General",
+          "shader": "dft-general.wgsl.jinja",
+          "derive": { "hasAxisInput": false, "hasDftLengthInput": true, "skipRuntimeAxis1": false },
+          "bindings": ["input", "output", "dft_len_in", "params_10"],
           "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.output) / dim(shapes.output, ranks.output - 1)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.output) / dim(shapes.output, ranks.output - 1)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
           }
         }
       ]

build/webgpu/metadata.json CHANGED Viewed

@@ -1,23 +1,42 @@
 {
   "name": "ai.onnx.DFT",
-  "id": "_ai_onnx_dft_webgpu_dbca2e5",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "jbPgs835f9m0aqqbSqY7zJJ+sRdHAdkjUzasXhjUEMM=",
-      "dft-contiguous-naive.wgsl.jinja": "HizdEBeuURKgl0Zxsz3xs62LJA6qmXvTYeYgC/XW5FE=",
-      "dft-general.wgsl.jinja": "fRvFH2YPz8DQLaTWSwXfScJErAFyeim5yqU39h+DYIc=",
-      "dft-rank4.wgsl.jinja": "QYtIZb6gpPc4fPIglkSt50Nr3KgbmuqoZk3EsuCgS3w=",
-      "dft-tiled-real.wgsl.jinja": "hbc3jiXVFIDZzBLo4O5hKQDqPIrNFdNktE85DhjWAzY=",
-      "fft-radix2-dit-storage.wgsl.jinja": "xLm5AhxAvxBdznovw1gQsQsaJGxusSmdtcaxUs00i2w=",
-      "fft-stockham-shared.wgsl.jinja": "7J4RLO2Q8puehryHRRHVhrrNxjNik5718G/qu6zoF6I=",
-      "manifest.json": "9IvgHvMofuhCPOTfi61UqnxP5X0IdE11o+JDI1V+dSs=",
-      "test.json": "pmZjTxGLBuze2IB7hjXRl0ks8BfzGeJ1OWCsaNDm/ao="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.DFT" }
 }

 {
   "name": "ai.onnx.DFT",
+  "id": "_ai_onnx_dft_webgpu_568fc49",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "9LPWvdkDFFGTQxPP69GFp4Z0lS287DSsZ2wxuGciXgY=",
+      "dft-contiguous-naive.wgsl.jinja": "K15oJXeOCAbiVr0x6U2+1TPOyj0PF+XNoFySxMFA4Do=",
+      "dft-general.wgsl.jinja": "pYttCZq/3h30wY4Vc6w4dEJJbEGvkbiHZqRxtXHLP7E=",
+      "dft-rank4.wgsl.jinja": "d4PWfZuv7gtxBLxVu7Wm7sH/MOsS7ZnoKpV1vcHqu1A=",
+      "dft-runtime-axis-fft-shared.wgsl.jinja": "6umWYtDkA3ebbfshjFFhmOVuWsygSpHu7vxGcdLd1Yc=",
+      "dft-tiled-real.wgsl.jinja": "mGmAT2TYrf1PWr+Y5Ebb33AMc/WT6Xr1/7icjk2feo0=",
+      "fft-radix2-dit-storage.wgsl.jinja": "oNrqAz6ZO92qD2wadfbTW006CaScjjvw3y6Hxp1+ydI=",
+      "fft-stockham-shared.wgsl.jinja": "J/1t+Z+A7lI2kgu0Om0U8+JNhrIQbCZOj2J9Fcdxd78=",
+      "manifest.json": "okwMs167d8CX1HETnuai4MaDVz17Sq8v1sKLmYwthOw=",
+      "test.json": "6+aX7hB91x77ThRU2xfezaBt56pn5UKH10Lano/9zag="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "runtime_axis_shared_fft": ["dft-general.wgsl.jinja", "dft-runtime-axis-fft-shared.wgsl.jinja"],
+      "batch_runtime_axis1_fft": ["dft-general.wgsl.jinja", "fft-stockham-shared.wgsl.jinja"],
+      "batch_axis1_fft_storage_onesided": ["fft-radix2-dit-storage.wgsl.jinja"],
+      "batch_axis1_fft_storage_inplace": ["fft-radix2-dit-storage.wgsl.jinja"],
+      "batch_axis1_tiled_real_nonpow2": ["dft-tiled-real.wgsl.jinja"],
+      "axis_fft_shared": ["fft-stockham-shared.wgsl.jinja"],
+      "batch_axis1_fft_explicit_equal_length": ["fft-stockham-shared.wgsl.jinja"],
+      "rank2_axis0": ["dft-contiguous-naive.wgsl.jinja"],
+      "batch_axis1": ["dft-contiguous-naive.wgsl.jinja"],
+      "rank4_axis2": ["dft-rank4.wgsl.jinja"],
+      "general_attr": ["dft-general.wgsl.jinja"],
+      "general_axis_len": ["dft-general.wgsl.jinja"],
+      "general_axis": ["dft-general.wgsl.jinja"],
+      "general_len": ["dft-general.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.DFT",
   "fixtureArrays": {
     "ort_rank2_complex_axis0_output_output": [36, 0, -4, 9.65685, -4, 4, -4, 1.65685, -4, 0, -4, -1.65685, -4, -4, -4, -9.65685]
   },
@@ -109,7 +108,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/signal/signal_ops_test.cc",
         "test": "SignalOpsTest.DFT17_Float_radix2",
-        "notes": "Complex companion to the real impulse edge: a subnormal complex impulse at sample zero should appear unchanged in every frequency bin."
       },
       "attrs": { "inverse": 0, "onesided": 0 },
       "inputs": {
@@ -130,7 +129,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/signal/signal_ops_test.cc",
         "test": "SignalOpsTest.DFT17_Float_inverse",
-        "notes": "Inverse companion: a subnormal DC bin scaled by 1/N should produce a valid finite subnormal constant signal."
       },
       "attrs": { "inverse": 1, "onesided": 0 },
       "inputs": {
@@ -151,7 +150,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/signal/signal_ops_test.cc",
         "test": "SignalOpsTest.DFT17_Float_radix2",
-        "notes": "Rank-2 axis-0 FFT companion: a real subnormal impulse should preserve the same subnormal real value in every bin."
       },
       "attrs": { "inverse": 0, "onesided": 0 },
       "inputs": {
@@ -168,7 +167,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/signal/signal_ops_test.cc",
         "test": "SignalOpsTest.DFT17_Float_radix2",
-        "notes": "Rank-4 axis-1 FFT companion: the subnormal impulse is along the first signal axis with one inner position."
       },
       "attrs": { "inverse": 0, "onesided": 0 },
       "inputs": {
@@ -190,7 +189,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/signal/signal_ops_test.cc",
         "test": "SignalOpsTest.DFT17_Float_radix2",
-        "notes": "Rank-4 axis-2 FFT companion: the subnormal impulse is along the second signal axis with one outer channel."
       },
       "attrs": { "inverse": 0, "onesided": 0 },
       "inputs": {
@@ -223,7 +222,7 @@
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/signal/signal_ops_test.cc",
         "test": "SignalOpsTest.DFT20_invertible_real",
-        "notes": "Opset-20 axis is an input tensor. This same-shape axis=0 case is a regression lock ensuring the runtime axis overrides the default axis attribute."
       },
       "inputs": {
         "input": {
@@ -330,7 +329,7 @@
     {
       "name": "fft_batch_axis1_len256_batch4_real_onesided_perf_compact",
       "provenance": {
-        "notes": "Compact correctness sibling for the batch-axis onesided FFT benchmark; preserves real input, axis=1, power-of-two length, and onesided output without benchmark-scale tensors."
       },
       "attrs": { "onesided": 1 },
       "inputs": { "input": { "dtype": "float32", "shape": [4, 256, 1] } },
@@ -1181,7 +1180,7 @@
     {
       "name": "general_len_rank4_axis1_dftlength_zeropad_len6",
       "provenance": {
-        "notes": "dft_length=6 zero-pads a rank-4 axis-1 signal of length 4 up to 6 samples, forcing the general_len variant (rank-4 coordinate decode + zero-pad guard n>=input_axis_length). Untested rank-4 dft_length path."
       },
       "attrs": { "inverse": 0, "onesided": 0 },
       "inputs": {
@@ -1219,7 +1218,7 @@
     {
       "name": "general_axis_len_negative_axis_rank3_dftlength_truncate",
       "provenance": {
-        "notes": "Negative axis input (-2 -> axis 1 for rank 3) combined with a truncating dft_length (6 -> 4) forces general_axis_len and exercises runtime resolve_axis() negative normalization together with signal truncation. Untested combination."
       },
       "attrs": { "inverse": 0, "onesided": 0 },
       "inputs": {
@@ -1248,7 +1247,7 @@
     {
       "name": "fft_explicit_equal_length128_batch4_onesided",
       "provenance": {
-        "notes": "Correctness lock for the optimized explicit-dft_length route when the runtime length equals the power-of-two signal length."
       },
       "attrs": { "inverse": 0, "onesided": 1 },
       "inputs": {
@@ -1307,7 +1306,7 @@
     {
       "name": "tiled_real_nonpow2_len4097_onesided_above_old_boundary",
       "provenance": {
-        "notes": "Proves the tiled direct transform is bounded by device workgroup storage rather than the former 4095 shape cutoff."
       },
       "attrs": { "onesided": 1 },
       "inputs": {
@@ -1337,7 +1336,7 @@
     {
       "name": "rank6_axis2_batch_and_inner",
       "provenance": {
-        "notes": "A rank the packing has never seen, with batch dims on both sides of the transform axis (outer 4, inner 2). If any rank ceiling had merely been raised rather than removed, this is the case that finds it."
       },
       "attrs": { "inverse": 0, "onesided": 0 },
       "inputs": {
@@ -1369,6 +1368,109 @@
           "data": { "kind": "values", "values": [3.0, 0.0, -1.0, 0.0] }
         }
       }
     }
   ]
 }

 {
   "fixtureArrays": {
     "ort_rank2_complex_axis0_output_output": [36, 0, -4, 9.65685, -4, 4, -4, 1.65685, -4, 0, -4, -1.65685, -4, -4, -4, -9.65685]
   },
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/signal/signal_ops_test.cc",
         "test": "SignalOpsTest.DFT17_Float_radix2",
+        "notes": "A subnormal complex impulse at sample zero must appear unchanged in every frequency bin."
       },
       "attrs": { "inverse": 0, "onesided": 0 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/signal/signal_ops_test.cc",
         "test": "SignalOpsTest.DFT17_Float_inverse",
+        "notes": "An inverse transform of a subnormal DC bin, scaled by 1/N, must produce a finite subnormal constant signal."
       },
       "attrs": { "inverse": 1, "onesided": 0 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/signal/signal_ops_test.cc",
         "test": "SignalOpsTest.DFT17_Float_radix2",
+        "notes": "A rank-2 axis-0 FFT of a real subnormal impulse must preserve that real value in every bin."
       },
       "attrs": { "inverse": 0, "onesided": 0 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/signal/signal_ops_test.cc",
         "test": "SignalOpsTest.DFT17_Float_radix2",
+        "notes": "A rank-4 axis-1 FFT places a subnormal impulse along the first signal axis with one inner position."
       },
       "attrs": { "inverse": 0, "onesided": 0 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/signal/signal_ops_test.cc",
         "test": "SignalOpsTest.DFT17_Float_radix2",
+        "notes": "A rank-4 axis-2 FFT places a subnormal impulse along the second signal axis with one outer channel."
       },
       "attrs": { "inverse": 0, "onesided": 0 },
       "inputs": {
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/signal/signal_ops_test.cc",
         "test": "SignalOpsTest.DFT20_invertible_real",
+        "notes": "In opset 20 the axis is an input tensor; axis=0 must override the operator's default axis on this shape."
       },
       "inputs": {
         "input": {
     {
       "name": "fft_batch_axis1_len256_batch4_real_onesided_perf_compact",
       "provenance": {
+        "notes": "A compact real-input FFT along batch axis 1 uses a power-of-two length and one-sided output without benchmark-scale tensors."
       },
       "attrs": { "onesided": 1 },
       "inputs": { "input": { "dtype": "float32", "shape": [4, 256, 1] } },
     {
       "name": "general_len_rank4_axis1_dftlength_zeropad_len6",
       "provenance": {
+        "notes": "dft_length=6 zero-pads a rank-4 axis-1 signal from four to six samples, exercising rank-4 coordinate decoding and the zero-pad boundary."
       },
       "attrs": { "inverse": 0, "onesided": 0 },
       "inputs": {
     {
       "name": "general_axis_len_negative_axis_rank3_dftlength_truncate",
       "provenance": {
+        "notes": "For rank 3, axis=-2 resolves to axis 1; dft_length=4 then truncates the six-sample signal on the general runtime-axis path."
       },
       "attrs": { "inverse": 0, "onesided": 0 },
       "inputs": {
     {
       "name": "fft_explicit_equal_length128_batch4_onesided",
       "provenance": {
+        "notes": "An explicit dft_length equal to the 128-sample signal length exercises the optimized power-of-two path without padding or truncation."
       },
       "attrs": { "inverse": 0, "onesided": 1 },
       "inputs": {
     {
       "name": "tiled_real_nonpow2_len4097_onesided_above_old_boundary",
       "provenance": {
+        "notes": "A 4,097-sample non-power-of-two transform exercises tiled direct DFT at the workgroup-storage eligibility boundary."
       },
       "attrs": { "onesided": 1 },
       "inputs": {
     {
       "name": "rank6_axis2_batch_and_inner",
       "provenance": {
+        "notes": "A rank-6 input with batch dimensions on both sides of the transform axis exercises coordinate decoding with outer size 4 and inner size 2."
       },
       "attrs": { "inverse": 0, "onesided": 0 },
       "inputs": {
           "data": { "kind": "values", "values": [3.0, 0.0, -1.0, 0.0] }
         }
       }
+    },
+    {
+      "name": "runtime_axis_fft_rank4_axis1_real_len64_inner3",
+      "attrs": {},
+      "inputs": {
+        "axis": { "dtype": "int32", "shape": [], "data": { "kind": "values", "values": [1] } },
+        "input": {
+          "dtype": "float32",
+          "shape": [2, 64, 3, 1],
+          "data": { "kind": "fillFloat32", "sinStep": 0.027, "cosStep": 0.033, "scale": 1.0 }
+        }
+      },
+      "outputs": {
+        "output": { "dtype": "float32", "shape": [2, 64, 3, 2], "tolerance": 0.005, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "runtime_axis_fft_rank4_axis0_complex_inverse_len128",
+      "attrs": { "inverse": 1 },
+      "inputs": {
+        "axis": { "dtype": "int32", "shape": [], "data": { "kind": "values", "values": [0] } },
+        "input": {
+          "dtype": "float32",
+          "shape": [128, 2, 3, 2],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.17, "scale": 3.0 }
+        }
+      },
+      "outputs": {
+        "output": { "dtype": "float32", "shape": [128, 2, 3, 2], "tolerance": 0.005, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "runtime_axis_fft_rank2_axis0_real_len256",
+      "attrs": {},
+      "inputs": {
+        "axis": { "dtype": "int32", "shape": [], "data": { "kind": "values", "values": [0] } },
+        "input": {
+          "dtype": "float32",
+          "shape": [256, 1],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.041, "scale": 1.0 }
+        }
+      },
+      "outputs": { "output": { "dtype": "float32", "shape": [256, 2], "tolerance": 0.01, "relTolerance": 0.0001 } }
+    },
+    {
+      "name": "runtime_axis_fft_rank5_axis2_len64_batch_and_inner",
+      "attrs": {},
+      "inputs": {
+        "axis": { "dtype": "int32", "shape": [], "data": { "kind": "values", "values": [2] } },
+        "input": {
+          "dtype": "float32",
+          "shape": [2, 2, 64, 3, 1],
+          "data": { "kind": "fillFloat32", "sinStep": 0.031, "cosStep": 0.023, "scale": 1.0 }
+        }
+      },
+      "outputs": {
+        "output": { "dtype": "float32", "shape": [2, 2, 64, 3, 2], "tolerance": 0.005, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "runtime_axis_general_len10_beside_fft_candidate_axis0",
+      "attrs": {},
+      "inputs": {
+        "axis": { "dtype": "int32", "shape": [], "data": { "kind": "values", "values": [1] } },
+        "input": {
+          "dtype": "float32",
+          "shape": [64, 10, 2, 1],
+          "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.02, "scale": 1.0 }
+        }
+      },
+      "outputs": {
+        "output": { "dtype": "float32", "shape": [64, 10, 2, 2], "tolerance": 0.002, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "runtime_axis_fft_shorter_candidate_axis2_len64_beside_len1024",
+      "attrs": {},
+      "inputs": {
+        "axis": { "dtype": "int32", "shape": [], "data": { "kind": "values", "values": [2] } },
+        "input": {
+          "dtype": "float32",
+          "shape": [2, 1024, 64, 1],
+          "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 1.0 }
+        }
+      },
+      "outputs": {
+        "output": { "dtype": "float32", "shape": [2, 1024, 64, 2], "tolerance": 0.005, "relTolerance": 0.0001 }
+      }
+    },
+    {
+      "name": "runtime_axis_fft_negative_axis_rank4_len256_complex",
+      "attrs": {},
+      "inputs": {
+        "axis": { "dtype": "int32", "shape": [], "data": { "kind": "values", "values": [-3] } },
+        "input": {
+          "dtype": "float32",
+          "shape": [3, 256, 2, 2],
+          "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.037, "scale": 2.0 }
+        }
+      },
+      "outputs": {
+        "output": { "dtype": "float32", "shape": [3, 256, 2, 2], "tolerance": 0.02, "relTolerance": 0.0001 }
+      }
     }
   ]
 }