1const FFT_RESOLUTION: u32 = 256u;
2const FFT_STAGES: u32 = 8u;
3
4@group(0) @binding(0) var source_field: texture_2d_array<f32>;
5@group(0) @binding(1) var target_field: texture_storage_2d_array<rgba32float, write>;
6
7var<workgroup> bank_a: array<vec4<f32>, FFT_RESOLUTION>;
8var<workgroup> bank_b: array<vec4<f32>, FFT_RESOLUTION>;
9
10fn multiply_pair(value: vec4<f32>, twiddle: vec2<f32>) -> vec4<f32> {
11    return vec4(
12        value.x * twiddle.x - value.y * twiddle.y,
13        value.x * twiddle.y + value.y * twiddle.x,
14        value.z * twiddle.x - value.w * twiddle.y,
15        value.z * twiddle.y + value.w * twiddle.x,
16    );
17}
18
19fn butterfly(index: u32, stage: u32, a: vec4<f32>, b: vec4<f32>) -> vec4<f32> {
20    let half_span = 1u << stage;
21    let span = 2u * half_span;
22    let j = index % half_span;
23    let angle = 2.0 * 3.141592653589793 * f32(j) / f32(span);
24    let rotated = multiply_pair(b, vec2(cos(angle), sin(angle)));
25    return select(a + rotated, a - rotated, (index % span) >= half_span);
26}
27
28fn transform(local_index: u32) {
29    for (var stage = 0u; stage < FFT_STAGES; stage += 1u) {
30        let half_span = 1u << stage;
31        let span = 2u * half_span;
32        let j = local_index % half_span;
33        let base = (local_index / span) * span + j;
34        if stage % 2u == 0u {
35            bank_b[local_index] = butterfly(local_index, stage, bank_a[base], bank_a[base + half_span]);
36        } else {
37            bank_a[local_index] = butterfly(local_index, stage, bank_b[base], bank_b[base + half_span]);
38        }
39        workgroupBarrier();
40    }
41}
42
43#ifdef FFT_VERTICAL
44@compute @workgroup_size(FFT_RESOLUTION, 1, 1)
45fn main(
46    @builtin(local_invocation_id) local: vec3<u32>,
47    @builtin(workgroup_id) group: vec3<u32>,
48) {
49    let index = local.x;
50    let source = vec2<i32>(i32(group.y), i32(reverseBits(index) >> 24u));
51    bank_a[index] = textureLoad(source_field, source, i32(group.z), 0);
52    workgroupBarrier();
53    transform(index);
54    textureStore(target_field, vec2<i32>(i32(group.y), i32(index)), i32(group.z), bank_a[index]);
55}
56#endif
57
58#ifndef FFT_VERTICAL
59@compute @workgroup_size(FFT_RESOLUTION, 1, 1)
60fn main(
61    @builtin(local_invocation_id) local: vec3<u32>,
62    @builtin(workgroup_id) group: vec3<u32>,
63) {
64    let index = local.x;
65    let source = vec2<i32>(i32(reverseBits(index) >> 24u), i32(group.y));
66    bank_a[index] = textureLoad(source_field, source, i32(group.z), 0);
67    workgroupBarrier();
68    transform(index);
69    textureStore(target_field, vec2<i32>(i32(index), i32(group.y)), i32(group.z), bank_a[index]);
70}
71#endif