Fourier stage: real-field transform, and radix-4 Stockham - concept-collection/turing-sphere

concept-collection / turing-sphere

Fourier stage: real-field transform, and radix-4 Stockham

Two independent changes to the FFT, both gated so the previous kernels stay reachable and testable. Real-field transform (SHT_REAL_FFT=0 to disable). The spatial field is real -- layout.ts stores m >= 0 only and implies Q_{l,-m} = (-1)^m conj(Q_lm) -- so a length-N transform is an N/2-point complex FFT plus a recombination, not a full N-point complex FFT of a Hermitian spectrum. With H = N/2: synthesis Z[k] = (X[k] + conj(X[H-k])) + i e^{+2pi i k/N} (X[k] - conj(X[H-k])) z = FFT_H^{+}(Z), x[2m] = Re z[m], x[2m+1] = Im z[m] analysis z[m] = x[2m] + i x[2m+1], Z = FFT_H^{-}(z) X[k] = (Z[k]+conj(Z[H-k]))/2 + e^{-2pi i k/N} * -i(Z[k]-conj(Z[H-k]))/2 The complex kernels are untouched: they are what a complex-valued spatial field would need, and stockham()/stockham4() are shared by both. Half the workgroup storage also matters for reach -- the fft path is gated on 16*nphi <= maxComputeWorkgroupStorageSize, which the real path effectively halves to 8*nphi. Radix-4 Stockham (SHT_RADIX=2 for the old one). log4(n) stages rather than log2(n), each stage being one workgroupBarrier. A leading radix-2 stage handles odd log2(n), so H=128 costs 1+3 stages instead of 7. Measured on an RTX PRO 6000 Blackwell, per-kernel GPU time from timestamp queries (scripts/_ts.ts), us: complex r2 real r2 real r4 128x256 fourSynth 6.40 5.95 5.95 fourAnalys 6.30 6.18 5.50 256x512 fourSynth 7.42 7.14 7.01 fourAnalys 7.36 7.04 6.30 Both are smaller than they look on paper: halving the arithmetic bought ~5% and halving the barriers ~0-11%, so the Fourier stage is bound by neither. It is most likely bound by memory: fm is laid out [m][ilat], which suits the Legendre kernels (a thread per latitude, coalesced) but not these, which own one latitude and read every m -- fm[k*NLAT + ilat] puts consecutive threads NLAT*8 bytes apart. Fixing that wants several latitudes per FFT workgroup so the reads coalesce along latitude; the halved shared memory here makes room for it. Not attempted. npm run bench --preset allencahn: lmax=127 11158 -> 12330 steps/s, lmax=255 6925 -> 7306. Against the original code, 5795 -> 12330 (2.13x) and 3246 -> 7306 (2.25x). diagnose-sht verifies both stages against the f64 reference on all four combinations of the two flags; test:node passes.

danfortunato <dan.fortunato@gmail.com> committed 20 days ago commit 9cb1d1af47b3 parent eb4a9e7 Browse files

2 changed files+255−7

src/sht/sht.tsmodified+60−4View file

@@ -13,6 +13,8 @@ import { legSynthWGSL, legAnalysWGSL } from './wgsl/leg.ts';

1313 import {

1414 fftSynthWGSL,

1515 fftAnalysWGSL,

16+ fftSynthRealWGSL,

17+ fftAnalysRealWGSL,

1618 dftSynthWGSL,

1719 dftAnalysWGSL,

1820 fftThreads,

@@ -238,18 +240,27 @@ export class ShtPlan {

238240 parity,

239241 };

240242 (this as { legLat: number }).legLat = parity ? nlat / 2 : nlat;

241- const fourP = { mmax, nlat, nphi };

243+ const fourP = { mmax, nlat, nphi, radix: (tuning('SHT_RADIX') as number | undefined) ?? 4 };

244+ // The spatial field is real (layout.ts stores m >= 0 only), so the Fourier

245+ // stage can run an nphi/2-point complex FFT plus a recombination instead of

246+ // a full nphi-point one: half the arithmetic and half the workgroup storage.

247+ // The complex kernels remain for a future complex-valued field, and are what

248+ // SHT_REAL_FFT=0 selects.

249+ const realFft =

250+ this.fourierMode === 'fft' && nphi % 2 === 0 && tuning('SHT_REAL_FFT') !== false;

251+ const fftS = realFft ? fftSynthRealWGSL : fftSynthWGSL;

252+ const fftA = realFft ? fftAnalysRealWGSL : fftAnalysWGSL;

242253 const [pLegS, pLegA, pFourS, pFourA] = await Promise.all([

243254 makePipeline(dev, legSynthWGSL(legP), 'leg_synth'),

244255 makePipeline(dev, legAnalysWGSL(legP), 'leg_analys'),

245256 makePipeline(

246257 dev,

247- this.fourierMode === 'fft' ? fftSynthWGSL(fourP) : dftSynthWGSL(fourP),

258+ this.fourierMode === 'fft' ? fftS(fourP) : dftSynthWGSL(fourP),

248259 this.fourierMode === 'fft' ? 'fft_synth' : 'dft_synth',

249260 ),

250261 makePipeline(

251262 dev,

252- this.fourierMode === 'fft' ? fftAnalysWGSL(fourP) : dftAnalysWGSL(fourP),

263+ this.fourierMode === 'fft' ? fftA(fourP) : dftAnalysWGSL(fourP),

253264 this.fourierMode === 'fft' ? 'fft_analys' : 'dft_analys',

254265 ),

255266 ]);

@@ -348,6 +359,47 @@ export class ShtPlan {

348359 pass.end();

349360 }

350361

362+ /**

363+ * Diagnostics: encode one stage alone, in its own pass, so a timestamp query

364+ * can measure just that kernel. The solver wants both stages in a shared pass

365+ * and should use encodeSynthInto/encodeAnalysInto; this exists because

366+ * inferring per-kernel cost by subtracting trivially-sized runs is unreliable.

367+ */

368+ encodeStage(

369+ encoder: GPUCommandEncoder,

370+ stage: 'legSynth' | 'fourSynth' | 'fourAnalys' | 'legAnalys',

371+ timestampWrites?: GPUComputePassTimestampWrites,

372+ ): void {

373+ const { mmax, nlat, nphi } = this.cfg;

374+ const fft = this.fourierMode === 'fft';

375+ const pass = encoder.beginComputePass({ label: `sht-${stage}`, timestampWrites });

376+ switch (stage) {

377+ case 'legSynth':

378+ pass.setPipeline(this.pipeLegSynth);

379+ pass.setBindGroup(0, this.bgLegSynth);

380+ pass.dispatchWorkgroups(Math.ceil(this.legLat / WG_SYNTH), mmax + 1);

381+ break;

382+ case 'fourSynth':

383+ pass.setPipeline(this.pipeFourSynth);

384+ pass.setBindGroup(0, this.bgFourSynth);

385+ if (fft) pass.dispatchWorkgroups(nlat);

386+ else pass.dispatchWorkgroups(Math.ceil(nphi / 64), nlat);

387+ break;

388+ case 'fourAnalys':

389+ pass.setPipeline(this.pipeFourAnalys);

390+ pass.setBindGroup(0, this.bgFourAnalys);

391+ if (fft) pass.dispatchWorkgroups(nlat);

392+ else pass.dispatchWorkgroups(Math.ceil((mmax + 1) / 64), nlat);

393+ break;

394+ case 'legAnalys':

395+ pass.setPipeline(this.pipeLegAnalys);

396+ pass.setBindGroup(0, this.bgLegAnalys);

397+ pass.dispatchWorkgroups(mmax + 1);

398+ break;

399+ }

400+ pass.end();

401+ }

402+

351403 /** Record the analysis (spatial spatBuf -> spectral qlmOut) into an encoder. */

352404 encodeAnalys(encoder: GPUCommandEncoder): void {

353405 const pass = encoder.beginComputePass({ label: 'sht-analys' });

@@ -429,7 +481,11 @@ export async function requestShtDevice(): Promise<GPUDevice> {

429481 // `subgroups` lets the analysis reduction use subgroupAdd instead of a

430482 // shared-memory tree (2 barriers per l-pair instead of 1 + log2(wgAnalys)).

431483 // Optional: ShtPlan falls back to the tree when it is not available.

432- const features: GPUFeatureName[] = adapter.features.has('subgroups') ? ['subgroups'] : [];

484+ const features: GPUFeatureName[] = [];

485+ if (adapter.features.has('subgroups')) features.push('subgroups');

486+ // timestamp-query is only used by the profiling scripts, but it has to be

487+ // requested at device creation, and asking costs nothing when unused.

488+ if (adapter.features.has('timestamp-query')) features.push('timestamp-query');

433489 return adapter.requestDevice({

434490 requiredFeatures: features,

435491 requiredLimits: { maxComputeWorkgroupStorageSize: wgStorage },

src/sht/wgsl/fourier.tsmodified+195−3View file

@@ -23,14 +23,22 @@ export interface FourierParams {

2323 mmax: number;

2424 nlat: number;

2525 nphi: number;

26+ /** 2 or 4; radix-4 uses log4(n) barrier stages instead of log2(n). */

27+ radix?: number;

2628 }

2729

2830 const TRIG_BINDING = /* wgsl */ `

2931 @group(0) @binding(2) var<storage, read> trig: array<vec2f>; // (cos,sin)(2*pi*k/NPHI), k < NPHI

3032 `;

3133

32-function stockham(nphi: number, threads: number, sign: number): string {

33- const log2n = Math.log2(nphi);

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+ */

39+function stockham(n: number, threads: number, sign: number, scale = 1): string {

40+ const nphi = n;

41+ const log2n = Math.log2(n);

3442 if (!Number.isInteger(log2n)) throw new Error('fft requires power-of-two nphi');

3543 // twiddle for pass with half-block ns: w = e^{sign*i*pi*j/ns} = T[j * (N/(2*ns))]^sign

3644 return /* wgsl */ `

@@ -56,7 +64,7 @@ fn fft_inplace(lid: u32) {

5664 workgroupBarrier();

5765 let ns = 1u << p;

5866 let sel = p & 1u;

59- let stride = ${nphi / 2}u >> p; // N/(2*ns)

67+ let stride = ${(nphi / 2) * scale}u >> p; // (NPHI/n) * n/(2*ns)

6068 for (var t = lid; t < ${nphi / 2}u; t += ${threads}u) {

6169 let j = t & (ns - 1u);

6270 let tw = trig[j * stride];

@@ -74,6 +82,93 @@ const FFT_OUT_SEL: u32 = ${log2n % 2}u;

7482 `;

7583 }

7684

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+ */

95+function 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 */ `

103+var<workgroup> bufA: array<vec2f, ${n}>;

104+var<workgroup> bufB: array<vec2f, ${n}>;

105+

106+fn 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+}

109+fn ld(sel: u32, i: u32) -> vec2f {

110+ if (sel == 0u) { return bufA[i]; }

111+ return bufB[i];

112+}

113+fn st_(sel: u32, i: u32, v: vec2f) {

114+ if (sel == 0u) { bufA[i] = v; } else { bufB[i] = v; }

115+}

116+fn tw(i: u32) -> vec2f {

117+ let t = trig[i];

118+ return vec2f(t.x, ${negY}t.y);

119+}

120+

121+fn 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+}

168+const FFT_OUT_SEL: u32 = ${total % 2}u;

169+`;

170+}

171+

77172 /** Choose FFT workgroup size: enough threads for the butterflies, capped at 256. */

78173 export function fftThreads(nphi: number): number {

79174 return Math.max(32, Math.min(256, nphi / 2));

@@ -192,3 +287,100 @@ fn dft_analys(@builtin(global_invocation_id) gid: vec3u) {

192287 }

193288 `;

194289 }

290+

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+ */

309+export function fftSynthRealWGSL(p: FourierParams): string {

310+ const H = p.nphi / 2;

311+ const T = fftThreads(H);

312+ return /* wgsl */ `

313+const MMAX: u32 = ${p.mmax}u;

314+const NLAT: u32 = ${p.nlat}u;

315+const NPHI: u32 = ${p.nphi}u;

316+const 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)}

321+

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.

324+fn 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+}

329+

330+@compute @workgroup_size(${T})

331+fn 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+}

351+

352+export function fftAnalysRealWGSL(p: FourierParams): string {

353+ const H = p.nphi / 2;

354+ const T = fftThreads(H);

355+ return /* wgsl */ `

356+const MMAX: u32 = ${p.mmax}u;

357+const NLAT: u32 = ${p.nlat}u;

358+const NPHI: u32 = ${p.nphi}u;

359+const 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)}

364+

365+@compute @workgroup_size(${T})

366+fn 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+}