yum/vfft

git clone https://git.yummers.dev/yum/vfft

yumFix shader names90039aa

master
branchesmasterall refs →
1.0 KiB42 linesraw
1import lib.util;
2
3[[vk::binding(0, 0)]]
4StructuredBuffer<float2> input;
5
6[[vk::binding(1, 0)]]
7RWStructuredBuffer<float2> output;
8
9[[vk::binding(2, 0)]]
10StructuredBuffer<float2> twiddles;
11
12[[vk::binding(3, 0)]]
13cbuffer Uniforms {
14  // Invariant: stride * big_n = input_len
15  uint stride;
16  uint big_n;
17  uint input_len;
18};
19
20// 256 threads per FFT
21[shader("compute")]
22[numthreads(256, 1, 1)]
23void fft_v1_batched(uint3 globalId : SV_DispatchThreadID)
24{
25    const uint id = globalId.x;
26    const uint stride = Uniforms.stride;
27    const uint big_n = Uniforms.big_n;
28    const uint half_big_n = big_n / 2;
29    const uint k = id / stride;
30
31    const uint which_fft = globalId.y;
32    const uint start_idx = which_fft * Uniforms.input_len + id % stride;
33
34    // Get odd and even elements.
35    let p = input[start_idx + 2 * k * stride];
36    let q = input[start_idx + (2 * k + 1) * stride];
37    // Combine.
38    let factor = twiddles[k * stride];
39    let qf = cmul(q, factor);
40    output[start_idx + k * stride] = p + qf;
41    output[start_idx + (k + half_big_n) * stride] = p - qf;
42}