1/**
2 * WGSL Fourier-stage kernels (the role cuFFT/VkFFT plays in SHTNS).
3 *
4 * Real fields, band-limited to |m| <= mmax < nphi/2:
5 * - synthesis: assemble a Hermitian spectrum from F_m (m >= 0) and do an
6 * inverse complex FFT along phi; take the real part.
7 * - analysis: forward complex FFT of the (real) row; keep m = 0..mmax.
8 *
9 * Two implementations, selected at plan creation:
10 * - 'fft': radix-2 Stockham in workgroup memory, one workgroup per
11 * latitude row. Requires nphi a power of two and
12 * 2 * 8 * nphi bytes <= maxComputeWorkgroupStorageSize.
13 * - 'dft': direct band-limited trigonometric summation, O(nphi * mmax)
14 * per row. Works for any nphi; also useful as a cross-check.
15 *
16 * All trigonometric factors come from a host-precomputed (f64 -> f32)
17 * table trig[k] = (cos, sin)(2*pi*k/nphi): device sin/cos is only
18 * guaranteed to ~2^-11 absolute error under Vulkan, which would dominate
19 * the fp32 transform error.
20 */
22export interface FourierParams {
23 mmax: number;
24 nlat: number;
25 nphi: number;
26 /** 2 or 4; radix-4 uses log4(n) barrier stages instead of log2(n). */
27 radix?: number;
28}
30const TRIG_BINDING = /* wgsl */ `
31@group(0) @binding(2) var<storage, read> trig: array<vec2f>; // (cos,sin)(2*pi*k/NPHI), k < NPHI
32`;
34/**
35 * @param n transform length (bufA/bufB are this long)
36 * @param scale trig-table stride multiplier: the table holds
37 * (cos,sin)(2*pi*k/NPHI), so an n-point transform needs NPHI/n.
38 */
39function stockham(n: number, threads: number, sign: number, scale = 1): string {
40 const nphi = n;
41 const log2n = Math.log2(n);
42 if (!Number.isInteger(log2n)) throw new Error('fft requires power-of-two nphi');
43 // twiddle for pass with half-block ns: w = e^{sign*i*pi*j/ns} = T[j * (N/(2*ns))]^sign
44 return /* wgsl */ `
45var<workgroup> bufA: array<vec2f, ${nphi}>;
46var<workgroup> bufB: array<vec2f, ${nphi}>;
48fn cmul(a: vec2f, b: vec2f) -> vec2f {
49 return vec2f(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
50}
52fn ld(sel: u32, i: u32) -> vec2f {
53 if (sel == 0u) { return bufA[i]; }
54 return bufB[i];
55}
56fn st_(sel: u32, i: u32, v: vec2f) {
57 if (sel == 0u) { bufA[i] = v; } else { bufB[i] = v; }
58}
60// radix-2 Stockham, natural order in and out; data starts in bufA (sel 0)
61// and ends in sel = LOG2N % 2. Unnormalized: X_k = sum_j x_j e^{s*2*pi*i*jk/N}.
62fn fft_inplace(lid: u32) {
63 for (var p = 0u; p < ${log2n}u; p++) {
64 workgroupBarrier();
65 let ns = 1u << p;
66 let sel = p & 1u;
67 let stride = ${(nphi / 2) * scale}u >> p; // (NPHI/n) * n/(2*ns)
68 for (var t = lid; t < ${nphi / 2}u; t += ${threads}u) {
69 let j = t & (ns - 1u);
70 let tw = trig[j * stride];
71 let w = vec2f(tw.x, ${sign > 0 ? '' : '-'}tw.y);
72 let u = ld(sel, t);
73 let v = cmul(ld(sel, t + ${nphi / 2}u), w);
74 let idst = 2u * (t - j) + j;
75 st_(1u - sel, idst, u + v);
76 st_(1u - sel, idst + ns, u - v);
77 }
78 }
79 workgroupBarrier();
80}
81const FFT_OUT_SEL: u32 = ${log2n % 2}u;
82`;
83}
85/**
86 * Radix-4 Stockham. Same interface and conventions as stockham(), but log4(n)
87 * stages instead of log2(n) -- each stage carries a workgroupBarrier, and
88 * barriers are what these kernels are actually bound by. When log2(n) is odd a
89 * single radix-2 stage runs first, so n = 128 costs 1 + 3 stages rather than 7.
90 *
91 * Butterfly, with w = e^{s 2 pi i / 4}:
92 * a = x0 + x2, b = x0 - x2, c = x1 + x3, d = s i (x1 - x3)
93 * y = (a + c, b + d, a - c, b - d)
94 */
95function stockham4(n: number, threads: number, sign: number, scale = 1): string {
96 const log2n = Math.log2(n);
97 if (!Number.isInteger(log2n)) throw new Error('fft requires power-of-two nphi');
98 const needR2 = log2n % 2 === 1;
99 const stages4 = Math.floor(log2n / 2);
100 const total = (needR2 ? 1 : 0) + stages4;
101 const negY = sign > 0 ? '' : '-';
102 return /* wgsl */ `
103var<workgroup> bufA: array<vec2f, ${n}>;
104var<workgroup> bufB: array<vec2f, ${n}>;
106fn cmul(a: vec2f, b: vec2f) -> vec2f {
107 return vec2f(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
108}
109fn ld(sel: u32, i: u32) -> vec2f {
110 if (sel == 0u) { return bufA[i]; }
111 return bufB[i];
112}
113fn st_(sel: u32, i: u32, v: vec2f) {
114 if (sel == 0u) { bufA[i] = v; } else { bufB[i] = v; }
115}
116fn tw(i: u32) -> vec2f {
117 let t = trig[i];
118 return vec2f(t.x, ${negY}t.y);
119}
121fn fft_inplace(lid: u32) {
122 var sel = 0u;
123 var ns = 1u;
124${
125 needR2
126 ? ` // leading radix-2 (ns = 1, so the twiddle is 1 and is skipped)
127 workgroupBarrier();
128 for (var t = lid; t < ${n / 2}u; t += ${threads}u) {
129 let u = ld(sel, t);
130 let v = ld(sel, t + ${n / 2}u);
131 st_(1u - sel, 2u * t, u + v);
132 st_(1u - sel, 2u * t + 1u, u - v);
133 }
134 sel = 1u - sel;
135 ns = 2u;`
136 : ''
137}
138 for (var p = 0u; p < ${stages4}u; p++) {
139 workgroupBarrier();
140 let s4 = ${(n * scale) / 4}u / ns; // trig unit: NPHI / (4 * ns)
141 for (var t = lid; t < ${n / 4}u; t += ${threads}u) {
142 let j = t & (ns - 1u);
143 let x0 = ld(sel, t);
144 var x1 = ld(sel, t + ${n / 4}u);
145 var x2 = ld(sel, t + ${n / 2}u);
146 var x3 = ld(sel, t + ${(3 * n) / 4}u);
147 if (ns > 1u) {
148 x1 = cmul(x1, tw(j * s4));
149 x2 = cmul(x2, tw(2u * j * s4));
150 x3 = cmul(x3, tw(3u * j * s4));
151 }
152 let a = x0 + x2;
153 let b = x0 - x2;
154 let c = x1 + x3;
155 let e = x1 - x3;
156 let d = vec2f(${sign > 0 ? '-e.y, e.x' : 'e.y, -e.x'}); // s * i * e
157 let idst = 4u * (t - j) + j;
158 st_(1u - sel, idst, a + c);
159 st_(1u - sel, idst + ns, b + d);
160 st_(1u - sel, idst + 2u * ns, a - c);
161 st_(1u - sel, idst + 3u * ns, b - d);
162 }
163 sel = 1u - sel;
164 ns = ns * 4u;
165 }
166 workgroupBarrier();
167}
168const FFT_OUT_SEL: u32 = ${total % 2}u;
169`;
170}
172/** Choose FFT workgroup size: enough threads for the butterflies, capped at 256. */
173export function fftThreads(nphi: number): number {
174 return Math.max(32, Math.min(256, nphi / 2));
175}
177export function fftSynthWGSL(p: FourierParams): string {
178 const T = fftThreads(p.nphi);
179 return /* wgsl */ `
180const MMAX: u32 = ${p.mmax}u;
181const NLAT: u32 = ${p.nlat}u;
182const NPHI: u32 = ${p.nphi}u;
183@group(0) @binding(0) var<storage, read> fm: array<vec2f>;
184@group(0) @binding(1) var<storage, read_write> spat: array<f32>;
185${TRIG_BINDING}
186${stockham(p.nphi, T, +1)}
188@compute @workgroup_size(${T})
189fn fft_synth(@builtin(local_invocation_id) lid3: vec3u,
190 @builtin(workgroup_id) wid: vec3u) {
191 let lid = lid3.x;
192 let ilat = wid.x;
193 // assemble Hermitian spectrum: X[0] = Re F_0, X[m] = F_m, X[N-m] = conj(F_m)
194 for (var k = lid; k < NPHI; k += ${T}u) {
195 var v = vec2f(0.0);
196 if (k == 0u) {
197 v = vec2f(fm[ilat].x, 0.0);
198 } else if (k <= MMAX) {
199 v = fm[k * NLAT + ilat];
200 } else if (k >= NPHI - MMAX) {
201 let c = fm[(NPHI - k) * NLAT + ilat];
202 v = vec2f(c.x, -c.y);
203 }
204 bufA[k] = v;
205 }
206 fft_inplace(lid);
207 for (var k = lid; k < NPHI; k += ${T}u) {
208 spat[ilat * NPHI + k] = ld(FFT_OUT_SEL, k).x;
209 }
210}
211`;
212}
214export function fftAnalysWGSL(p: FourierParams): string {
215 const T = fftThreads(p.nphi);
216 return /* wgsl */ `
217const MMAX: u32 = ${p.mmax}u;
218const NLAT: u32 = ${p.nlat}u;
219const NPHI: u32 = ${p.nphi}u;
220@group(0) @binding(0) var<storage, read> spat: array<f32>;
221@group(0) @binding(1) var<storage, read_write> fm: array<vec2f>;
222${TRIG_BINDING}
223${stockham(p.nphi, T, -1)}
225@compute @workgroup_size(${T})
226fn fft_analys(@builtin(local_invocation_id) lid3: vec3u,
227 @builtin(workgroup_id) wid: vec3u) {
228 let lid = lid3.x;
229 let ilat = wid.x;
230 for (var k = lid; k < NPHI; k += ${T}u) {
231 bufA[k] = vec2f(spat[ilat * NPHI + k], 0.0);
232 }
233 fft_inplace(lid);
234 for (var m = lid; m <= MMAX; m += ${T}u) {
235 fm[m * NLAT + ilat] = ld(FFT_OUT_SEL, m);
236 }
237}
238`;
239}
241export function dftSynthWGSL(p: FourierParams): string {
242 return /* wgsl */ `
243const MMAX: u32 = ${p.mmax}u;
244const NLAT: u32 = ${p.nlat}u;
245const NPHI: u32 = ${p.nphi}u;
246@group(0) @binding(0) var<storage, read> fm: array<vec2f>;
247@group(0) @binding(1) var<storage, read_write> spat: array<f32>;
248${TRIG_BINDING}
250@compute @workgroup_size(64)
251fn dft_synth(@builtin(global_invocation_id) gid: vec3u) {
252 let iphi = gid.x;
253 let ilat = gid.y;
254 if (iphi >= NPHI) { return; }
255 var v: f32 = fm[ilat].x; // m = 0: real part
256 for (var m = 1u; m <= MMAX; m++) {
257 let w = trig[(m * iphi) % NPHI]; // e^{+i m phi}
258 let c = fm[m * NLAT + ilat];
259 v += 2.0 * (c.x * w.x - c.y * w.y);
260 }
261 spat[ilat * NPHI + iphi] = v;
262}
263`;
264}
266export function dftAnalysWGSL(p: FourierParams): string {
267 return /* wgsl */ `
268const MMAX: u32 = ${p.mmax}u;
269const NLAT: u32 = ${p.nlat}u;
270const NPHI: u32 = ${p.nphi}u;
271@group(0) @binding(0) var<storage, read> spat: array<f32>;
272@group(0) @binding(1) var<storage, read_write> fm: array<vec2f>;
273${TRIG_BINDING}
275@compute @workgroup_size(64)
276fn dft_analys(@builtin(global_invocation_id) gid: vec3u) {
277 let m = gid.x;
278 let ilat = gid.y;
279 if (m > MMAX) { return; }
280 var acc = vec2f(0.0);
281 for (var j = 0u; j < NPHI; j++) {
282 let w = trig[(m * j) % NPHI]; // conj => e^{-i m phi}
283 let f = spat[ilat * NPHI + j];
284 acc += f * vec2f(w.x, -w.y);
285 }
286 fm[m * NLAT + ilat] = acc;
287}
288`;
289}
291/**
292 * Real-field Fourier stage: half the arithmetic and half the shared memory of
293 * the complex path, which transforms N points to get a Hermitian result.
294 *
295 * A length-N real transform is an N/2-point complex FFT wrapped in a
296 * recombination. Writing H = N/2 and taking the unnormalized conventions of the
297 * complex kernels above (synthesis e^{+i}, analysis e^{-i}):
298 *
299 * synthesis Z[k] = (X[k] + conj(X[H-k])) + i e^{+2pi i k/N} (X[k] - conj(X[H-k]))
300 * z = FFT_H^{+}(Z), then x[2m] = Re z[m], x[2m+1] = Im z[m]
301 * analysis z[m] = x[2m] + i x[2m+1], Z = FFT_H^{-}(z)
302 * Xe = (Z[k] + conj(Z[H-k]))/2, Xo = -i (Z[k] - conj(Z[H-k]))/2
303 * X[k] = Xe + e^{-2pi i k/N} Xo
304 *
305 * The factors of 2 in the synthesis direction cancel against the 1/2 in Xe/Xo,
306 * which is why none appear there. The complex kernels are kept: they are what a
307 * complex-valued spatial field would use, and stockham() is shared by both.
308 */
309export function fftSynthRealWGSL(p: FourierParams): string {
310 const H = p.nphi / 2;
311 const T = fftThreads(H);
312 return /* wgsl */ `
313const MMAX: u32 = ${p.mmax}u;
314const NLAT: u32 = ${p.nlat}u;
315const NPHI: u32 = ${p.nphi}u;
316const H: u32 = ${H}u;
317@group(0) @binding(0) var<storage, read> fm: array<vec2f>;
318@group(0) @binding(1) var<storage, read_write> spat: array<f32>;
319${TRIG_BINDING}
320${(p.radix ?? 4) === 4 ? stockham4(H, T, +1, p.nphi / H) : stockham(H, T, +1, p.nphi / H)}
322// X[k] of the Hermitian spectrum, for 0 <= k <= H. mmax < H, so the
323// upper-conjugate branch of the complex kernel cannot be reached here.
324fn spec(ilat: u32, k: u32) -> vec2f {
325 if (k == 0u) { return vec2f(fm[ilat].x, 0.0); }
326 if (k <= MMAX) { return fm[k * NLAT + ilat]; }
327 return vec2f(0.0);
328}
330@compute @workgroup_size(${T})
331fn fft_synth(@builtin(local_invocation_id) lid3: vec3u,
332 @builtin(workgroup_id) wid: vec3u) {
333 let lid = lid3.x;
334 let ilat = wid.x;
335 for (var k = lid; k < H; k += ${T}u) {
336 let xk = spec(ilat, k);
337 let xh = spec(ilat, H - k);
338 let cj = vec2f(xh.x, -xh.y);
339 let b = cmul(xk - cj, trig[k]); // e^{+2 pi i k / N}
340 bufA[k] = (xk + cj) + vec2f(-b.y, b.x); // + i * b
341 }
342 fft_inplace(lid);
343 for (var m = lid; m < H; m += ${T}u) {
344 let z = ld(FFT_OUT_SEL, m);
345 spat[ilat * NPHI + 2u * m] = z.x;
346 spat[ilat * NPHI + 2u * m + 1u] = z.y;
347 }
348}
349`;
350}
352export function fftAnalysRealWGSL(p: FourierParams): string {
353 const H = p.nphi / 2;
354 const T = fftThreads(H);
355 return /* wgsl */ `
356const MMAX: u32 = ${p.mmax}u;
357const NLAT: u32 = ${p.nlat}u;
358const NPHI: u32 = ${p.nphi}u;
359const H: u32 = ${H}u;
360@group(0) @binding(0) var<storage, read> spat: array<f32>;
361@group(0) @binding(1) var<storage, read_write> fm: array<vec2f>;
362${TRIG_BINDING}
363${(p.radix ?? 4) === 4 ? stockham4(H, T, -1, p.nphi / H) : stockham(H, T, -1, p.nphi / H)}
365@compute @workgroup_size(${T})
366fn fft_analys(@builtin(local_invocation_id) lid3: vec3u,
367 @builtin(workgroup_id) wid: vec3u) {
368 let lid = lid3.x;
369 let ilat = wid.x;
370 for (var m = lid; m < H; m += ${T}u) {
371 bufA[m] = vec2f(spat[ilat * NPHI + 2u * m], spat[ilat * NPHI + 2u * m + 1u]);
372 }
373 fft_inplace(lid);
374 for (var k = lid; k <= MMAX; k += ${T}u) {
375 let zk = ld(FFT_OUT_SEL, k);
376 let zh = ld(FFT_OUT_SEL, (H - k) % H); // Z[H] == Z[0]
377 let cj = vec2f(zh.x, -zh.y);
378 let xe = 0.5 * (zk + cj);
379 let d = 0.5 * (zk - cj);
380 let xo = vec2f(d.y, -d.x); // -i * d
381 let w = vec2f(trig[k].x, -trig[k].y); // e^{-2 pi i k / N}
382 fm[k * NLAT + ilat] = xe + cmul(xo, w);
383 }
384}
385`;
386}