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;
38const WG_ANALYS = 256;
40async function makePipeline(
41 device: GPUDevice,
42 code: string,
43 entryPoint: string,
44): Promise<GPUComputePipeline> {
45 device.pushErrorScope('validation');
46 const module = device.createShaderModule({ code, label: entryPoint });
47 const info = await module.getCompilationInfo();
48 const errors = info.messages.filter((m) => m.type === 'error');
49 if (errors.length) {
50 throw new Error(
51 `WGSL compile error in ${entryPoint}:\n` +
52 errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n'),
53 );
54 }
55 const pipeline = await device.createComputePipelineAsync({
56 layout: 'auto',
57 compute: { module, entryPoint },
58 label: entryPoint,
59 });
60 const err = await device.popErrorScope();
61 if (err) throw new Error(`pipeline ${entryPoint}: ${err.message}`);
62 return pipeline;
63}
65export class ShtPlan {
66 readonly cfg: ShtConfig;
67 readonly nlm: number;
68 readonly fourierMode: 'fft' | 'dft';
69 /** Colatitudes theta_i (f64, increasing: north to south). */
70 readonly theta: Float64Array;
71 readonly cosTheta: Float64Array;
72 readonly gaussWeights: Float64Array;
74 private device: GPUDevice;
75 private bufAb!: GPUBuffer;
76 private bufAmm!: GPUBuffer;
77 private bufCtstw!: GPUBuffer;
78 private bufTrig!: GPUBuffer;
79 /** Spectral input (synthesis) — write with queue.writeBuffer or use synth(). */
80 readonly qlmIn!: GPUBuffer;
81 /** Spectral output (analysis). */
82 readonly qlmOut!: GPUBuffer;
83 /** Fourier-space intermediate [(m)*nlat + ilat], complex f32. COPY_SRC so the
84 * stage boundary is observable: a transform is Legendre-then-Fourier, and
85 * scripts/diagnose-sht.ts tells the two apart by reading this. */
86 readonly fmBuf!: GPUBuffer;
87 /** Spatial field [ilat*nphi + iphi], f32. */
88 readonly spatBuf!: GPUBuffer;
89 private stageSpat!: GPUBuffer;
90 private stageQ!: GPUBuffer;
92 private pipeLegSynth!: GPUComputePipeline;
93 private pipeLegAnalys!: GPUComputePipeline;
94 private pipeFourSynth!: GPUComputePipeline;
95 private pipeFourAnalys!: GPUComputePipeline;
96 private bgLegSynth!: GPUBindGroup;
97 private bgLegAnalys!: GPUBindGroup;
98 private bgFourSynth!: GPUBindGroup;
99 private bgFourAnalys!: GPUBindGroup;
101 private constructor(device: GPUDevice, cfg: ShtConfig, fourierMode: 'fft' | 'dft') {
102 this.device = device;
103 this.cfg = cfg;
104 this.nlm = nlmCalc(cfg.lmax, cfg.mmax);
105 this.fourierMode = fourierMode;
106 const { x, w } = gaussNodesWeights(cfg.nlat);
107 this.cosTheta = x;
108 this.gaussWeights = w;
109 this.theta = new Float64Array(cfg.nlat);
110 for (let i = 0; i < cfg.nlat; i++) this.theta[i] = Math.acos(x[i]);
111 }
113 static async create(device: GPUDevice, cfg: ShtConfig, opts: ShtOptions = {}): Promise<ShtPlan> {
114 validateConfig(cfg);
115 const want = opts.fourier ?? 'auto';
116 const fftFits =
117 isPowerOfTwo(cfg.nphi) &&
118 16 * cfg.nphi <= device.limits.maxComputeWorkgroupStorageSize &&
119 fftThreads(cfg.nphi) <= device.limits.maxComputeInvocationsPerWorkgroup;
120 if (want === 'fft' && !fftFits) {
121 throw new Error(
122 `fourier:'fft' requires power-of-two nphi with 16*nphi <= maxComputeWorkgroupStorageSize ` +
123 `(nphi=${cfg.nphi}, limit=${device.limits.maxComputeWorkgroupStorageSize})`,
124 );
125 }
126 const mode: 'fft' | 'dft' = want === 'dft' ? 'dft' : fftFits ? 'fft' : 'dft';
127 const plan = new ShtPlan(device, cfg, mode);
128 await plan.init();
129 return plan;
130 }
132 private async init(): Promise<void> {
133 const { lmax, mmax, nlat, nphi } = this.cfg;
134 const dev = this.device;
135 const self = this as {
136 -readonly [k in keyof ShtPlan]: ShtPlan[k];
137 };
139 // --- host precomputation (f64), then downcast to f32 for upload ---
140 const { amm, ab } = legendreCoeffs(lmax, mmax);
141 const ctstw = new Float32Array(3 * nlat);
142 for (let i = 0; i < nlat; i++) {
143 ctstw[i] = this.cosTheta[i];
144 ctstw[nlat + i] = Math.sqrt(1 - this.cosTheta[i] * this.cosTheta[i]);
145 ctstw[2 * nlat + i] = this.gaussWeights[i] * ((2 * Math.PI) / nphi);
146 }
147 // twiddle/phase table in f64 (device sin/cos is too inaccurate: ~2^-11 under Vulkan)
148 const trig = new Float32Array(2 * nphi);
149 for (let k = 0; k < nphi; k++) {
150 trig[2 * k] = Math.cos((2 * Math.PI * k) / nphi);
151 trig[2 * k + 1] = Math.sin((2 * Math.PI * k) / nphi);
152 }
154 const mkBuf = (label: string, size: number, usage: GPUBufferUsageFlags) =>
155 dev.createBuffer({ label, size, usage });
156 this.bufAb = mkBuf('sht-ab', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
157 this.bufAmm = mkBuf('sht-amm', 4 * (mmax + 1), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
158 this.bufCtstw = mkBuf('sht-ctstw', 4 * 3 * nlat, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
159 this.bufTrig = mkBuf('sht-trig', 8 * nphi, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
160 self.qlmIn = mkBuf('sht-qlm-in', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC);
161 self.qlmOut = mkBuf('sht-qlm-out', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
162 self.fmBuf = mkBuf('sht-fm', 8 * (mmax + 1) * nlat, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
163 self.spatBuf = mkBuf('sht-spat', 4 * nlat * nphi, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC);
164 this.stageSpat = mkBuf('sht-stage-spat', 4 * nlat * nphi, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
165 this.stageQ = mkBuf('sht-stage-q', 8 * this.nlm, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
167 dev.queue.writeBuffer(this.bufAb, 0, new Float32Array(ab));
168 dev.queue.writeBuffer(this.bufAmm, 0, new Float32Array(amm));
169 dev.queue.writeBuffer(this.bufCtstw, 0, ctstw);
170 dev.queue.writeBuffer(this.bufTrig, 0, trig);
172 // --- shaders / pipelines ---
173 const legP = { lmax, mmax, nlat, wgSynth: WG_SYNTH, wgAnalys: WG_ANALYS };
174 const fourP = { mmax, nlat, nphi };
175 const [pLegS, pLegA, pFourS, pFourA] = await Promise.all([
176 makePipeline(dev, legSynthWGSL(legP), 'leg_synth'),
177 makePipeline(dev, legAnalysWGSL(legP), 'leg_analys'),
178 makePipeline(
179 dev,
180 this.fourierMode === 'fft' ? fftSynthWGSL(fourP) : dftSynthWGSL(fourP),
181 this.fourierMode === 'fft' ? 'fft_synth' : 'dft_synth',
182 ),
183 makePipeline(
184 dev,
185 this.fourierMode === 'fft' ? fftAnalysWGSL(fourP) : dftAnalysWGSL(fourP),
186 this.fourierMode === 'fft' ? 'fft_analys' : 'dft_analys',
187 ),
188 ]);
189 this.pipeLegSynth = pLegS;
190 this.pipeLegAnalys = pLegA;
191 this.pipeFourSynth = pFourS;
192 this.pipeFourAnalys = pFourA;
194 const entries = bgEntries;
195 this.bgLegSynth = dev.createBindGroup({
196 layout: pLegS.getBindGroupLayout(0),
197 entries: entries([this.bufAb, this.bufAmm, this.bufCtstw, this.qlmIn, this.fmBuf]),
198 });
199 this.bgLegAnalys = dev.createBindGroup({
200 layout: pLegA.getBindGroupLayout(0),
201 entries: entries([this.bufAb, this.bufAmm, this.bufCtstw, this.fmBuf, this.qlmOut]),
202 });
203 this.bgFourSynth = dev.createBindGroup({
204 layout: pFourS.getBindGroupLayout(0),
205 entries: entries([this.fmBuf, this.spatBuf, this.bufTrig]),
206 });
207 this.bgFourAnalys = dev.createBindGroup({
208 layout: pFourA.getBindGroupLayout(0),
209 entries: entries([this.spatBuf, this.fmBuf, this.bufTrig]),
210 });
211 }
213 /**
214 * Bind groups for one transform against caller-supplied spectral/spatial
215 * buffers, so a transform can read and write buffers it does not own (the
216 * .m-driven executor keeps a buffer per IR variable). Build these once at
217 * plan time, not per step. `fmBuf` stays internal scratch: passes and
218 * dispatches within a submission execute in order, so sequential transforms
219 * can share it.
220 */
221 createSynthBinding(qlmIn: GPUBuffer, spatOut: GPUBuffer): ShtBinding {
222 return {
223 bgLeg: this.device.createBindGroup({
224 layout: this.pipeLegSynth.getBindGroupLayout(0),
225 entries: bgEntries([this.bufAb, this.bufAmm, this.bufCtstw, qlmIn, this.fmBuf]),
226 }),
227 bgFour: this.device.createBindGroup({
228 layout: this.pipeFourSynth.getBindGroupLayout(0),
229 entries: bgEntries([this.fmBuf, spatOut, this.bufTrig]),
230 }),
231 };
232 }
234 createAnalysBinding(spatIn: GPUBuffer, qlmOut: GPUBuffer): ShtBinding {
235 return {
236 bgFour: this.device.createBindGroup({
237 layout: this.pipeFourAnalys.getBindGroupLayout(0),
238 entries: bgEntries([spatIn, this.fmBuf, this.bufTrig]),
239 }),
240 bgLeg: this.device.createBindGroup({
241 layout: this.pipeLegAnalys.getBindGroupLayout(0),
242 entries: bgEntries([this.bufAb, this.bufAmm, this.bufCtstw, this.fmBuf, qlmOut]),
243 }),
244 };
245 }
247 /** Record synthesis into an existing compute pass. */
248 encodeSynthInto(pass: GPUComputePassEncoder, b: ShtBinding): void {
249 const { mmax, nlat, nphi } = this.cfg;
250 pass.setPipeline(this.pipeLegSynth);
251 pass.setBindGroup(0, b.bgLeg);
252 pass.dispatchWorkgroups(Math.ceil(nlat / WG_SYNTH), mmax + 1);
253 pass.setPipeline(this.pipeFourSynth);
254 pass.setBindGroup(0, b.bgFour);
255 if (this.fourierMode === 'fft') {
256 pass.dispatchWorkgroups(nlat);
257 } else {
258 pass.dispatchWorkgroups(Math.ceil(nphi / 64), nlat);
259 }
260 }
262 /** Record analysis into an existing compute pass. */
263 encodeAnalysInto(pass: GPUComputePassEncoder, b: ShtBinding): void {
264 const { mmax, nlat } = this.cfg;
265 pass.setPipeline(this.pipeFourAnalys);
266 pass.setBindGroup(0, b.bgFour);
267 if (this.fourierMode === 'fft') {
268 pass.dispatchWorkgroups(nlat);
269 } else {
270 pass.dispatchWorkgroups(Math.ceil((mmax + 1) / 64), nlat);
271 }
272 pass.setPipeline(this.pipeLegAnalys);
273 pass.setBindGroup(0, b.bgLeg);
274 pass.dispatchWorkgroups(mmax + 1);
275 }
277 /** Record the synthesis (spectral qlmIn -> spatial spatBuf) into an encoder. */
278 encodeSynth(encoder: GPUCommandEncoder): void {
279 const pass = encoder.beginComputePass({ label: 'sht-synth' });
280 this.encodeSynthInto(pass, { bgLeg: this.bgLegSynth, bgFour: this.bgFourSynth });
281 pass.end();
282 }
284 /** Record the analysis (spatial spatBuf -> spectral qlmOut) into an encoder. */
285 encodeAnalys(encoder: GPUCommandEncoder): void {
286 const pass = encoder.beginComputePass({ label: 'sht-analys' });
287 this.encodeAnalysInto(pass, { bgLeg: this.bgLegAnalys, bgFour: this.bgFourAnalys });
288 pass.end();
289 }
291 /**
292 * Spectral -> spatial. qlm: interleaved [re, im], SHTNS LM ordering,
293 * length 2*nlm. Returns the spatial field, length nlat*nphi.
294 */
295 async synth(qlm: Float32Array): Promise<Float32Array> {
296 const { nlat, nphi } = this.cfg;
297 if (qlm.length !== 2 * this.nlm) throw new Error(`qlm must have length ${2 * this.nlm}`);
298 this.device.queue.writeBuffer(this.qlmIn, 0, qlm as Float32Array<ArrayBuffer>);
299 const enc = this.device.createCommandEncoder();
300 this.encodeSynth(enc);
301 enc.copyBufferToBuffer(this.spatBuf, 0, this.stageSpat, 0, 4 * nlat * nphi);
302 this.device.queue.submit([enc.finish()]);
303 await this.stageSpat.mapAsync(GPUMapMode.READ);
304 const out = new Float32Array(this.stageSpat.getMappedRange().slice(0));
305 this.stageSpat.unmap();
306 return out;
307 }
309 /** Spatial -> spectral. spat: length nlat*nphi. Returns interleaved qlm, length 2*nlm. */
310 async analys(spat: Float32Array): Promise<Float32Array> {
311 const { nlat, nphi } = this.cfg;
312 if (spat.length !== nlat * nphi) throw new Error(`spat must have length ${nlat * nphi}`);
313 this.device.queue.writeBuffer(this.spatBuf, 0, spat as Float32Array<ArrayBuffer>);
314 const enc = this.device.createCommandEncoder();
315 this.encodeAnalys(enc);
316 enc.copyBufferToBuffer(this.qlmOut, 0, this.stageQ, 0, 8 * this.nlm);
317 this.device.queue.submit([enc.finish()]);
318 await this.stageQ.mapAsync(GPUMapMode.READ);
319 const out = new Float32Array(this.stageQ.getMappedRange().slice(0));
320 this.stageQ.unmap();
321 return out;
322 }
324 destroy(): void {
325 for (const b of [
326 this.bufAb, this.bufAmm, this.bufCtstw, this.bufTrig, this.qlmIn, this.qlmOut,
327 this.fmBuf, this.spatBuf, this.stageSpat, this.stageQ,
328 ]) b?.destroy();
329 }
330}
332/** Best-effort human-readable adapter name, so it is clear which GPU (or
333 * software rasterizer) is actually running the transforms. */
334export async function describeAdapter(device: GPUDevice): Promise<string> {
335 const fmt = (info: GPUAdapterInfo | undefined): string => {
336 if (!info) return '';
337 const parts = [info.description, info.device, info.vendor].filter(
338 (s): s is string => !!s && s.length > 0,
339 );
340 const name = parts[0] ?? '';
341 return info.architecture && !name.includes(info.architecture)
342 ? `${name} (${info.architecture})`.trim()
343 : name;
344 };
345 const own = fmt((device as GPUDevice & { adapterInfo?: GPUAdapterInfo }).adapterInfo);
346 if (own) return own;
347 try {
348 const adapter = await navigator.gpu.requestAdapter();
349 return fmt(adapter?.info);
350 } catch {
351 return '';
352 }
353}
355/** Request an adapter/device suitable for the transforms. */
356export async function requestShtDevice(): Promise<GPUDevice> {
357 if (!navigator.gpu) throw new Error('WebGPU is not available in this browser');
358 const adapter = await navigator.gpu.requestAdapter();
359 if (!adapter) throw new Error('No WebGPU adapter available');
360 // ask for a larger workgroup storage if the adapter offers it (bigger FFTs)
361 const wgStorage = Math.min(adapter.limits.maxComputeWorkgroupStorageSize, 32768);
362 return adapter.requestDevice({
363 requiredLimits: { maxComputeWorkgroupStorageSize: wgStorage },
364 });
365}