1/**
2 * WebGPU spherical harmonic transform plan (scalar transforms, fp32).
3 *
4 * Mirrors the structure of the SHTNS CUDA backend (sht_gpu.cu):
5 * host-side f64 precomputation of grid + recurrence coefficients, shader
6 * source generated with sizes baked in (SHTNS uses NVRTC; WGSL is always
7 * runtime-compiled), then per-transform: Legendre stage + Fourier stage.
8 */
9import { gaussNodesWeights } from './gauss.ts';
10import { legendreCoeffs } from './coeffs.ts';
11import { nlmCalc, validateConfig, isPowerOfTwo, type ShtConfig } from './layout.ts';
12import { legSynthWGSL, legAnalysWGSL } from './wgsl/leg.ts';
13import {
14 fftSynthWGSL,
15 fftAnalysWGSL,
16 dftSynthWGSL,
17 dftAnalysWGSL,
18 fftThreads,
19} from './wgsl/fourier.ts';
21export type FourierMode = 'auto' | 'fft' | 'dft';
23/** The two bind groups (Legendre stage, Fourier stage) of one transform. */
24export interface ShtBinding {
25 readonly bgLeg: GPUBindGroup;
26 readonly bgFour: GPUBindGroup;
27}
29const bgEntries = (bufs: GPUBuffer[]) =>
30 bufs.map((buffer, binding) => ({ binding, resource: { buffer } }));
32export interface ShtOptions {
33 /** Fourier stage implementation. 'auto' picks fft when nphi is a power of two that fits in workgroup memory. */
34 fourier?: FourierMode;
35}
37const WG_SYNTH = 64;
39/**
40 * Tuning knob, for A/B-ing a change without editing code. Reads globalThis
41 * first (set it before creating a plan, as scripts/_ab.ts does), then the
42 * environment, so `SHT_SUBGROUPS=0 npm run bench:sht` works too. `process` is
43 * absent in the browser, where only the globalThis form applies.
44 */
45function tuning(name: string): unknown {
46 const g = (globalThis as Record<string, unknown>)[name];
47 if (g !== undefined) return g;
48 const env = (globalThis as { process?: { env?: Record<string, string | undefined> } }).process?.env?.[name];
49 if (env === undefined || env === '') return undefined;
50 if (env === '1' || env === 'true') return true;
51 if (env === '0' || env === 'false') return false;
52 const n = Number(env);
53 return Number.isFinite(n) ? n : env;
54}
56/**
57 * Workgroup size for the analysis Legendre reduction. The right answer differs
58 * between the two reduction strategies, so it is chosen per strategy.
59 *
60 * Measured on an RTX PRO 6000 Blackwell (analysis, us). Shared-memory tree,
61 * where each doubling of wgAnalys costs another barrier per l-pair:
62 *
63 * wgAnalys: 16 32 64 128 256
64 * nlat=128 43.2 40.6 42.6 46.7 52.9 -> 32
65 * nlat=256 104.0 85.6 87.8 89.9 100.0 -> 32
66 * nlat=512 408.8 216.9 184.4 190.8 204.8 -> 64
67 *
68 * i.e. max(32, nlat/8). A flat 32 would be worse than the old default of 256 at
69 * nlat=512, so it cannot be fitted on one grid. With subgroupAdd the barrier
70 * count stops growing with wgAnalys and the picture inverts: threads in flight,
71 * (mmax+1) * wgAnalys, becomes binding, since analysis dispatches only mmax+1
72 * workgroups. 128 then wins at every grid (round trip, us):
73 *
74 * 128x256 37.8 (vs 38.3), 256x512 66.6 (vs 74.7), 512x1024 131.9 (vs 156.6)
75 */
76function defaultWgAnalys(nlat: number, limit: number, subgroups: boolean): number {
77 if (subgroups) {
78 // capped at nlat so small grids do not launch threads with no latitude to own
79 let cap = 1;
80 while (cap < nlat) cap *= 2;
81 return Math.min(128, limit, cap);
82 }
83 const target = Math.max(32, nlat / 8);
84 let wg = 1;
85 while (wg < target) wg *= 2; // the tree reduction halves, so a power of two
86 return Math.min(wg, limit);
87}
89async function makePipeline(
90 device: GPUDevice,
91 code: string,
92 entryPoint: string,
93): Promise<GPUComputePipeline> {
94 device.pushErrorScope('validation');
95 const module = device.createShaderModule({ code, label: entryPoint });
96 const info = await module.getCompilationInfo();
97 const errors = info.messages.filter((m) => m.type === 'error');
98 if (errors.length) {
99 throw new Error(
100 `WGSL compile error in ${entryPoint}:\n` +
101 errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n'),
102 );
103 }
104 const pipeline = await device.createComputePipelineAsync({
105 layout: 'auto',
106 compute: { module, entryPoint },
107 label: entryPoint,
108 });
109 const err = await device.popErrorScope();
110 if (err) throw new Error(`pipeline ${entryPoint}: ${err.message}`);
111 return pipeline;
112}
114export class ShtPlan {
115 readonly cfg: ShtConfig;
116 readonly nlm: number;
117 readonly fourierMode: 'fft' | 'dft';
118 /** Colatitudes theta_i (f64, increasing: north to south). */
119 readonly theta: Float64Array;
120 readonly cosTheta: Float64Array;
121 readonly gaussWeights: Float64Array;
123 private device: GPUDevice;
124 private bufAb!: GPUBuffer;
125 private bufAmm!: GPUBuffer;
126 private bufCtstw!: GPUBuffer;
127 private bufTrig!: GPUBuffer;
128 /** Spectral input (synthesis) — write with queue.writeBuffer or use synth(). */
129 readonly qlmIn!: GPUBuffer;
130 /** Spectral output (analysis). */
131 readonly qlmOut!: GPUBuffer;
132 /** Fourier-space intermediate [(m)*nlat + ilat], complex f32. COPY_SRC so the
133 * stage boundary is observable: a transform is Legendre-then-Fourier, and
134 * scripts/diagnose-sht.ts tells the two apart by reading this. */
135 readonly fmBuf!: GPUBuffer;
136 /** Spatial field [ilat*nphi + iphi], f32. */
137 readonly spatBuf!: GPUBuffer;
138 private stageSpat!: GPUBuffer;
139 private stageQ!: GPUBuffer;
141 private pipeLegSynth!: GPUComputePipeline;
142 private pipeLegAnalys!: GPUComputePipeline;
143 private pipeFourSynth!: GPUComputePipeline;
144 private pipeFourAnalys!: GPUComputePipeline;
145 private bgLegSynth!: GPUBindGroup;
146 private bgLegAnalys!: GPUBindGroup;
147 private bgFourSynth!: GPUBindGroup;
148 private bgFourAnalys!: GPUBindGroup;
150 private constructor(device: GPUDevice, cfg: ShtConfig, fourierMode: 'fft' | 'dft') {
151 this.device = device;
152 this.cfg = cfg;
153 this.nlm = nlmCalc(cfg.lmax, cfg.mmax);
154 this.fourierMode = fourierMode;
155 const { x, w } = gaussNodesWeights(cfg.nlat);
156 this.cosTheta = x;
157 this.gaussWeights = w;
158 this.theta = new Float64Array(cfg.nlat);
159 for (let i = 0; i < cfg.nlat; i++) this.theta[i] = Math.acos(x[i]);
160 }
162 static async create(device: GPUDevice, cfg: ShtConfig, opts: ShtOptions = {}): Promise<ShtPlan> {
163 validateConfig(cfg);
164 const want = opts.fourier ?? 'auto';
165 const fftFits =
166 isPowerOfTwo(cfg.nphi) &&
167 16 * cfg.nphi <= device.limits.maxComputeWorkgroupStorageSize &&
168 fftThreads(cfg.nphi) <= device.limits.maxComputeInvocationsPerWorkgroup;
169 if (want === 'fft' && !fftFits) {
170 throw new Error(
171 `fourier:'fft' requires power-of-two nphi with 16*nphi <= maxComputeWorkgroupStorageSize ` +
172 `(nphi=${cfg.nphi}, limit=${device.limits.maxComputeWorkgroupStorageSize})`,
173 );
174 }
175 const mode: 'fft' | 'dft' = want === 'dft' ? 'dft' : fftFits ? 'fft' : 'dft';
176 const plan = new ShtPlan(device, cfg, mode);
177 await plan.init();
178 return plan;
179 }
181 private async init(): Promise<void> {
182 const { lmax, mmax, nlat, nphi } = this.cfg;
183 const dev = this.device;
184 const self = this as {
185 -readonly [k in keyof ShtPlan]: ShtPlan[k];
186 };
188 // --- host precomputation (f64), then downcast to f32 for upload ---
189 const { amm, ab } = legendreCoeffs(lmax, mmax);
190 const ctstw = new Float32Array(3 * nlat);
191 for (let i = 0; i < nlat; i++) {
192 ctstw[i] = this.cosTheta[i];
193 ctstw[nlat + i] = Math.sqrt(1 - this.cosTheta[i] * this.cosTheta[i]);
194 ctstw[2 * nlat + i] = this.gaussWeights[i] * ((2 * Math.PI) / nphi);
195 }
196 // twiddle/phase table in f64 (device sin/cos is too inaccurate: ~2^-11 under Vulkan)
197 const trig = new Float32Array(2 * nphi);
198 for (let k = 0; k < nphi; k++) {
199 trig[2 * k] = Math.cos((2 * Math.PI * k) / nphi);
200 trig[2 * k + 1] = Math.sin((2 * Math.PI * k) / nphi);
201 }
203 const mkBuf = (label: string, size: number, usage: GPUBufferUsageFlags) =>
204 dev.createBuffer({ label, size, usage });
205 this.bufAb = mkBuf('sht-ab', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
206 this.bufAmm = mkBuf('sht-amm', 4 * (mmax + 1), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
207 this.bufCtstw = mkBuf('sht-ctstw', 4 * 3 * nlat, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
208 this.bufTrig = mkBuf('sht-trig', 8 * nphi, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
209 self.qlmIn = mkBuf('sht-qlm-in', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC);
210 self.qlmOut = mkBuf('sht-qlm-out', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
211 self.fmBuf = mkBuf('sht-fm', 8 * (mmax + 1) * nlat, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
212 self.spatBuf = mkBuf('sht-spat', 4 * nlat * nphi, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC);
213 this.stageSpat = mkBuf('sht-stage-spat', 4 * nlat * nphi, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
214 this.stageQ = mkBuf('sht-stage-q', 8 * this.nlm, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
216 dev.queue.writeBuffer(this.bufAb, 0, new Float32Array(ab));
217 dev.queue.writeBuffer(this.bufAmm, 0, new Float32Array(amm));
218 dev.queue.writeBuffer(this.bufCtstw, 0, ctstw);
219 dev.queue.writeBuffer(this.bufTrig, 0, trig);
221 // --- shaders / pipelines ---
222 const subgroups = tuning('SHT_SUBGROUPS') !== false && dev.features.has('subgroups');
223 const wgAnalys =
224 (tuning('SHT_WG_ANALYS') as number | undefined) ??
225 defaultWgAnalys(nlat, dev.limits.maxComputeInvocationsPerWorkgroup, subgroups);
226 const legP = {
227 lmax,
228 mmax,
229 nlat,
230 wgSynth: WG_SYNTH,
231 wgAnalys,
232 subgroups,
233 spanPairs: tuning('SHT_SPAN_PAIRS') as number | undefined,
234 };
235 const fourP = { mmax, nlat, nphi };
236 const [pLegS, pLegA, pFourS, pFourA] = await Promise.all([
237 makePipeline(dev, legSynthWGSL(legP), 'leg_synth'),
238 makePipeline(dev, legAnalysWGSL(legP), 'leg_analys'),
239 makePipeline(
240 dev,
241 this.fourierMode === 'fft' ? fftSynthWGSL(fourP) : dftSynthWGSL(fourP),
242 this.fourierMode === 'fft' ? 'fft_synth' : 'dft_synth',
243 ),
244 makePipeline(
245 dev,
246 this.fourierMode === 'fft' ? fftAnalysWGSL(fourP) : dftAnalysWGSL(fourP),
247 this.fourierMode === 'fft' ? 'fft_analys' : 'dft_analys',
248 ),
249 ]);
250 this.pipeLegSynth = pLegS;
251 this.pipeLegAnalys = pLegA;
252 this.pipeFourSynth = pFourS;
253 this.pipeFourAnalys = pFourA;
255 const entries = bgEntries;
256 this.bgLegSynth = dev.createBindGroup({
257 layout: pLegS.getBindGroupLayout(0),
258 entries: entries([this.bufAb, this.bufAmm, this.bufCtstw, this.qlmIn, this.fmBuf]),
259 });
260 this.bgLegAnalys = dev.createBindGroup({
261 layout: pLegA.getBindGroupLayout(0),
262 entries: entries([this.bufAb, this.bufAmm, this.bufCtstw, this.fmBuf, this.qlmOut]),
263 });
264 this.bgFourSynth = dev.createBindGroup({
265 layout: pFourS.getBindGroupLayout(0),
266 entries: entries([this.fmBuf, this.spatBuf, this.bufTrig]),
267 });
268 this.bgFourAnalys = dev.createBindGroup({
269 layout: pFourA.getBindGroupLayout(0),
270 entries: entries([this.spatBuf, this.fmBuf, this.bufTrig]),
271 });
272 }
274 /**
275 * Bind groups for one transform against caller-supplied spectral/spatial
276 * buffers, so a transform can read and write buffers it does not own (the
277 * .m-driven executor keeps a buffer per IR variable). Build these once at
278 * plan time, not per step. `fmBuf` stays internal scratch: passes and
279 * dispatches within a submission execute in order, so sequential transforms
280 * can share it.
281 */
282 createSynthBinding(qlmIn: GPUBuffer, spatOut: GPUBuffer): ShtBinding {
283 return {
284 bgLeg: this.device.createBindGroup({
285 layout: this.pipeLegSynth.getBindGroupLayout(0),
286 entries: bgEntries([this.bufAb, this.bufAmm, this.bufCtstw, qlmIn, this.fmBuf]),
287 }),
288 bgFour: this.device.createBindGroup({
289 layout: this.pipeFourSynth.getBindGroupLayout(0),
290 entries: bgEntries([this.fmBuf, spatOut, this.bufTrig]),
291 }),
292 };
293 }
295 createAnalysBinding(spatIn: GPUBuffer, qlmOut: GPUBuffer): ShtBinding {
296 return {
297 bgFour: this.device.createBindGroup({
298 layout: this.pipeFourAnalys.getBindGroupLayout(0),
299 entries: bgEntries([spatIn, this.fmBuf, this.bufTrig]),
300 }),
301 bgLeg: this.device.createBindGroup({
302 layout: this.pipeLegAnalys.getBindGroupLayout(0),
303 entries: bgEntries([this.bufAb, this.bufAmm, this.bufCtstw, this.fmBuf, qlmOut]),
304 }),
305 };
306 }
308 /** Record synthesis into an existing compute pass. */
309 encodeSynthInto(pass: GPUComputePassEncoder, b: ShtBinding): void {
310 const { mmax, nlat, nphi } = this.cfg;
311 pass.setPipeline(this.pipeLegSynth);
312 pass.setBindGroup(0, b.bgLeg);
313 pass.dispatchWorkgroups(Math.ceil(nlat / WG_SYNTH), mmax + 1);
314 pass.setPipeline(this.pipeFourSynth);
315 pass.setBindGroup(0, b.bgFour);
316 if (this.fourierMode === 'fft') {
317 pass.dispatchWorkgroups(nlat);
318 } else {
319 pass.dispatchWorkgroups(Math.ceil(nphi / 64), nlat);
320 }
321 }
323 /** Record analysis into an existing compute pass. */
324 encodeAnalysInto(pass: GPUComputePassEncoder, b: ShtBinding): void {
325 const { mmax, nlat } = this.cfg;
326 pass.setPipeline(this.pipeFourAnalys);
327 pass.setBindGroup(0, b.bgFour);
328 if (this.fourierMode === 'fft') {
329 pass.dispatchWorkgroups(nlat);
330 } else {
331 pass.dispatchWorkgroups(Math.ceil((mmax + 1) / 64), nlat);
332 }
333 pass.setPipeline(this.pipeLegAnalys);
334 pass.setBindGroup(0, b.bgLeg);
335 pass.dispatchWorkgroups(mmax + 1);
336 }
338 /** Record the synthesis (spectral qlmIn -> spatial spatBuf) into an encoder. */
339 encodeSynth(encoder: GPUCommandEncoder): void {
340 const pass = encoder.beginComputePass({ label: 'sht-synth' });
341 this.encodeSynthInto(pass, { bgLeg: this.bgLegSynth, bgFour: this.bgFourSynth });
342 pass.end();
343 }
345 /** Record the analysis (spatial spatBuf -> spectral qlmOut) into an encoder. */
346 encodeAnalys(encoder: GPUCommandEncoder): void {
347 const pass = encoder.beginComputePass({ label: 'sht-analys' });
348 this.encodeAnalysInto(pass, { bgLeg: this.bgLegAnalys, bgFour: this.bgFourAnalys });
349 pass.end();
350 }
352 /**
353 * Spectral -> spatial. qlm: interleaved [re, im], SHTNS LM ordering,
354 * length 2*nlm. Returns the spatial field, length nlat*nphi.
355 */
356 async synth(qlm: Float32Array): Promise<Float32Array> {
357 const { nlat, nphi } = this.cfg;
358 if (qlm.length !== 2 * this.nlm) throw new Error(`qlm must have length ${2 * this.nlm}`);
359 this.device.queue.writeBuffer(this.qlmIn, 0, qlm as Float32Array<ArrayBuffer>);
360 const enc = this.device.createCommandEncoder();
361 this.encodeSynth(enc);
362 enc.copyBufferToBuffer(this.spatBuf, 0, this.stageSpat, 0, 4 * nlat * nphi);
363 this.device.queue.submit([enc.finish()]);
364 await this.stageSpat.mapAsync(GPUMapMode.READ);
365 const out = new Float32Array(this.stageSpat.getMappedRange().slice(0));
366 this.stageSpat.unmap();
367 return out;
368 }
370 /** Spatial -> spectral. spat: length nlat*nphi. Returns interleaved qlm, length 2*nlm. */
371 async analys(spat: Float32Array): Promise<Float32Array> {
372 const { nlat, nphi } = this.cfg;
373 if (spat.length !== nlat * nphi) throw new Error(`spat must have length ${nlat * nphi}`);
374 this.device.queue.writeBuffer(this.spatBuf, 0, spat as Float32Array<ArrayBuffer>);
375 const enc = this.device.createCommandEncoder();
376 this.encodeAnalys(enc);
377 enc.copyBufferToBuffer(this.qlmOut, 0, this.stageQ, 0, 8 * this.nlm);
378 this.device.queue.submit([enc.finish()]);
379 await this.stageQ.mapAsync(GPUMapMode.READ);
380 const out = new Float32Array(this.stageQ.getMappedRange().slice(0));
381 this.stageQ.unmap();
382 return out;
383 }
385 destroy(): void {
386 for (const b of [
387 this.bufAb, this.bufAmm, this.bufCtstw, this.bufTrig, this.qlmIn, this.qlmOut,
388 this.fmBuf, this.spatBuf, this.stageSpat, this.stageQ,
389 ]) b?.destroy();
390 }
391}
393/** Best-effort human-readable adapter name, so it is clear which GPU (or
394 * software rasterizer) is actually running the transforms. */
395export async function describeAdapter(device: GPUDevice): Promise<string> {
396 const fmt = (info: GPUAdapterInfo | undefined): string => {
397 if (!info) return '';
398 const parts = [info.description, info.device, info.vendor].filter(
399 (s): s is string => !!s && s.length > 0,
400 );
401 const name = parts[0] ?? '';
402 return info.architecture && !name.includes(info.architecture)
403 ? `${name} (${info.architecture})`.trim()
404 : name;
405 };
406 const own = fmt((device as GPUDevice & { adapterInfo?: GPUAdapterInfo }).adapterInfo);
407 if (own) return own;
408 try {
409 const adapter = await navigator.gpu.requestAdapter();
410 return fmt(adapter?.info);
411 } catch {
412 return '';
413 }
414}
416/** Request an adapter/device suitable for the transforms. */
417export async function requestShtDevice(): Promise<GPUDevice> {
418 if (!navigator.gpu) throw new Error('WebGPU is not available in this browser');
419 const adapter = await navigator.gpu.requestAdapter();
420 if (!adapter) throw new Error('No WebGPU adapter available');
421 // ask for a larger workgroup storage if the adapter offers it (bigger FFTs)
422 const wgStorage = Math.min(adapter.limits.maxComputeWorkgroupStorageSize, 32768);
423 // `subgroups` lets the analysis reduction use subgroupAdd instead of a
424 // shared-memory tree (2 barriers per l-pair instead of 1 + log2(wgAnalys)).
425 // Optional: ShtPlan falls back to the tree when it is not available.
426 const features: GPUFeatureName[] = adapter.features.has('subgroups') ? ['subgroups'] : [];
427 return adapter.requestDevice({
428 requiredFeatures: features,
429 requiredLimits: { maxComputeWorkgroupStorageSize: wgStorage },
430 });
431}