yum/vfft

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

yumFix shader names90039aa

master
branchesmasterall refs →
2.2 KiB81 linesraw
1import lib.util;
2
3[[vk::binding(0, 0)]]
4RWStructuredBuffer<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  uint input_len;
15  uint n_iter;
16};
17
18groupshared float2 scratch[4096];
19
20void fft_v2_stage(uint which_fft, uint id, uint stride, uint big_n) {
21    const uint k = id / stride;
22    const uint start_idx = id % stride;
23
24    // Collect inputs.
25    let i0 = scratch[start_idx + 4 * k * stride];
26    let i1 = scratch[start_idx + (4 * k + 1) * stride];
27    let i2 = scratch[start_idx + (4 * k + 2) * stride];
28    let i3 = scratch[start_idx + (4 * k + 3) * stride];
29    // Collect relevant twiddles.
30    let ot1 = twiddles[1 * k * stride];
31    let ot2 = twiddles[2 * k * stride];
32    let ot3 = twiddles[3 * k * stride];
33
34    let a = i0;
35    let b = cmul(ot1, i1);
36    let c = cmul(ot2, i2);
37    let d = cmul(ot3, i3);
38
39    // To derive this, write the output assignments in terms of
40    // a/b/c/d, then factor out!
41    let ac_sum = a + c;
42    let ac_diff = a - c;
43    let bd_sum = b + d;
44    let bd_diff_ni = mul_ni(b - d);
45
46    GroupMemoryBarrierWithGroupSync();
47
48    scratch[start_idx + k * stride] = ac_sum + bd_sum;
49    scratch[start_idx + (k + big_n / 4) * stride] = ac_diff + bd_diff_ni;
50    scratch[start_idx + (k + big_n / 2) * stride] = ac_sum - bd_sum;
51    scratch[start_idx + (k + 3 * big_n / 4) * stride] = ac_diff - bd_diff_ni;
52
53    GroupMemoryBarrierWithGroupSync();
54}
55
56[shader("compute")]
57[numthreads(64, 16, 1)]
58void fft_v3_groupshared(uint3 globalId : SV_DispatchThreadID, uint id: SV_GroupIndex)
59{
60    const uint which_fft = globalId.z;
61
62    uint stride = Uniforms.input_len;
63    uint big_n = 1;
64
65    const uint offset = which_fft * Uniforms.input_len;
66    for (uint i = 0; i < 4; i++) {
67        scratch[id + i * 1024] = input[offset + id + i * 1024];
68    }
69    AllMemoryBarrierWithGroupSync();
70
71    for (uint i = 0; i < Uniforms.n_iter; i++) {
72      stride /= 4;
73      big_n *= 4;
74      fft_v2_stage(which_fft, id, stride, big_n);
75    }
76
77    for (uint i = 0; i < 4; i++) {
78        output[offset + id + i * 1024] = scratch[id + i * 1024];
79    }
80    AllMemoryBarrierWithGroupSync();
81}