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