yum/vfft

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

yumFix shader names90039aa

master
branchesmasterall refs →
3.2 KiB109 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 + 8 * k * stride];
26    let i1 = scratch[start_idx + (8 * k + 1) * stride];
27    let i2 = scratch[start_idx + (8 * k + 2) * stride];
28    let i3 = scratch[start_idx + (8 * k + 3) * stride];
29    let i4 = scratch[start_idx + (8 * k + 4) * stride];
30    let i5 = scratch[start_idx + (8 * k + 5) * stride];
31    let i6 = scratch[start_idx + (8 * k + 6) * stride];
32    let i7 = scratch[start_idx + (8 * k + 7) * stride];
33
34    // Collect relevant twiddles.
35    let ot1 = twiddles[1 * k * stride];
36    let ot2 = twiddles[2 * k * stride];
37    let ot3 = twiddles[3 * k * stride];
38    let ot4 = twiddles[4 * k * stride];
39    let ot5 = twiddles[5 * k * stride];
40    let ot6 = twiddles[6 * k * stride];
41    let ot7 = twiddles[7 * k * stride];
42
43    let a = i0;
44    let b = cmul(ot1, i1);
45    let c = cmul(ot2, i2);
46    let d = cmul(ot3, i3);
47    let e = cmul(ot4, i4);
48    let f = cmul(ot5, i5);
49    let g = cmul(ot6, i6);
50    let h = cmul(ot7, i7);
51
52    let ae_sum = a + e;
53    let ae_diff = a - e;
54    let bf_sum = b + f;
55    let bf_diff = b - f;
56    let cg_sum = c + g;
57    let cg_diff = c - g;
58    let dh_sum = d + h;
59    let dh_diff = d - h;
60
61    let w00 = ae_sum + cg_sum;
62    let w01 = ae_sum - cg_sum;
63    let w10 = ae_diff + rot_270(cg_diff);
64    let w11 = ae_diff - rot_270(cg_diff);
65    let x00 = bf_sum + dh_sum;
66    let x01 = rot_270(bf_sum) + rot_90(dh_sum);
67    let x10 = rot_45(rot_270(bf_diff) + rot_180(dh_diff));
68    let x11 = rot_45(rot_180(bf_diff) + rot_270(dh_diff));
69
70    GroupMemoryBarrierWithGroupSync();
71
72    scratch[start_idx + k * stride] = w00 + x00;
73    scratch[start_idx + (k + big_n / 8) * stride] = w10 + x10;
74    scratch[start_idx + (k + big_n / 4) * stride] = w01 + x01;
75    scratch[start_idx + (k + 3 * big_n / 8) * stride] = w11 + x11;
76    scratch[start_idx + (k + big_n / 2) * stride] = w00 - x00;
77    scratch[start_idx + (k + 5 * big_n / 8) * stride] = w10 - x10;
78    scratch[start_idx + (k + 3 * big_n / 4) * stride] = w01 - x01;
79    scratch[start_idx + (k + 7 * big_n / 8) * stride] = w11 - x11;
80
81    GroupMemoryBarrierWithGroupSync();
82}
83
84[shader("compute")]
85[numthreads(64, 8, 1)]
86void fft_v4_radix_8(uint3 globalId : SV_DispatchThreadID, uint id: SV_GroupIndex)
87{
88    const uint which_fft = globalId.z;
89
90    uint stride = Uniforms.input_len;
91    uint big_n = 1;
92
93    const uint offset = which_fft * Uniforms.input_len;
94    for (uint i = 0; i < 8; i++) {
95        scratch[id + i * 512] = input[offset + id + i * 512];
96    }
97    AllMemoryBarrierWithGroupSync();
98
99    for (uint i = 0; i < Uniforms.n_iter; i++) {
100      stride /= 8;
101      big_n *= 8;
102      fft_v2_stage(which_fft, id, stride, big_n);
103    }
104
105    for (uint i = 0; i < 8; i++) {
106        output[offset + id + i * 512] = scratch[id + i * 512];
107    }
108    AllMemoryBarrierWithGroupSync();
109}