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';
23export interface ShtOptions {
24 /** Fourier stage implementation. 'auto' picks fft when nphi is a power of two that fits in workgroup memory. */
25 fourier?: FourierMode;
26}
28const WG_SYNTH = 64;
29const WG_ANALYS = 256;
31async function makePipeline(
32 device: GPUDevice,
33 code: string,
34 entryPoint: string,
35): Promise<GPUComputePipeline> {
36 device.pushErrorScope('validation');
37 const module = device.createShaderModule({ code, label: entryPoint });
38 const info = await module.getCompilationInfo();
39 const errors = info.messages.filter((m) => m.type === 'error');
40 if (errors.length) {
41 throw new Error(
42 `WGSL compile error in ${entryPoint}:\n` +
43 errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n'),
44 );
45 }
46 const pipeline = await device.createComputePipelineAsync({
47 layout: 'auto',
48 compute: { module, entryPoint },
49 label: entryPoint,
50 });
51 const err = await device.popErrorScope();
52 if (err) throw new Error(`pipeline ${entryPoint}: ${err.message}`);
53 return pipeline;
54}
56export class ShtPlan {
57 readonly cfg: ShtConfig;
58 readonly nlm: number;
59 readonly fourierMode: 'fft' | 'dft';
60 /** Colatitudes theta_i (f64, increasing: north to south). */
61 readonly theta: Float64Array;
62 readonly cosTheta: Float64Array;
63 readonly gaussWeights: Float64Array;
65 private device: GPUDevice;
66 private bufAb!: GPUBuffer;
67 private bufAmm!: GPUBuffer;
68 private bufCtstw!: GPUBuffer;
69 private bufTrig!: GPUBuffer;
70 /** Spectral input (synthesis) — write with queue.writeBuffer or use synth(). */
71 readonly qlmIn!: GPUBuffer;
72 /** Spectral output (analysis). */
73 readonly qlmOut!: GPUBuffer;
74 /** Fourier-space intermediate [(m)*nlat + ilat], complex f32. */
75 readonly fmBuf!: GPUBuffer;
76 /** Spatial field [ilat*nphi + iphi], f32. */
77 readonly spatBuf!: GPUBuffer;
78 private stageSpat!: GPUBuffer;
79 private stageQ!: GPUBuffer;
81 private pipeLegSynth!: GPUComputePipeline;
82 private pipeLegAnalys!: GPUComputePipeline;
83 private pipeFourSynth!: GPUComputePipeline;
84 private pipeFourAnalys!: GPUComputePipeline;
85 private bgLegSynth!: GPUBindGroup;
86 private bgLegAnalys!: GPUBindGroup;
87 private bgFourSynth!: GPUBindGroup;
88 private bgFourAnalys!: GPUBindGroup;
90 private constructor(device: GPUDevice, cfg: ShtConfig, fourierMode: 'fft' | 'dft') {
91 this.device = device;
92 this.cfg = cfg;
93 this.nlm = nlmCalc(cfg.lmax, cfg.mmax);
94 this.fourierMode = fourierMode;
95 const { x, w } = gaussNodesWeights(cfg.nlat);
96 this.cosTheta = x;
97 this.gaussWeights = w;
98 this.theta = new Float64Array(cfg.nlat);
99 for (let i = 0; i < cfg.nlat; i++) this.theta[i] = Math.acos(x[i]);
100 }
102 static async create(device: GPUDevice, cfg: ShtConfig, opts: ShtOptions = {}): Promise<ShtPlan> {
103 validateConfig(cfg);
104 const want = opts.fourier ?? 'auto';
105 const fftFits =
106 isPowerOfTwo(cfg.nphi) &&
107 16 * cfg.nphi <= device.limits.maxComputeWorkgroupStorageSize &&
108 fftThreads(cfg.nphi) <= device.limits.maxComputeInvocationsPerWorkgroup;
109 if (want === 'fft' && !fftFits) {
110 throw new Error(
111 `fourier:'fft' requires power-of-two nphi with 16*nphi <= maxComputeWorkgroupStorageSize ` +
112 `(nphi=${cfg.nphi}, limit=${device.limits.maxComputeWorkgroupStorageSize})`,
113 );
114 }
115 const mode: 'fft' | 'dft' = want === 'dft' ? 'dft' : fftFits ? 'fft' : 'dft';
116 const plan = new ShtPlan(device, cfg, mode);
117 await plan.init();
118 return plan;
119 }
121 private async init(): Promise<void> {
122 const { lmax, mmax, nlat, nphi } = this.cfg;
123 const dev = this.device;
124 const self = this as {
125 -readonly [k in keyof ShtPlan]: ShtPlan[k];
126 };
128 // --- host precomputation (f64), then downcast to f32 for upload ---
129 const { amm, ab } = legendreCoeffs(lmax, mmax);
130 const ctstw = new Float32Array(3 * nlat);
131 for (let i = 0; i < nlat; i++) {
132 ctstw[i] = this.cosTheta[i];
133 ctstw[nlat + i] = Math.sqrt(1 - this.cosTheta[i] * this.cosTheta[i]);
134 ctstw[2 * nlat + i] = this.gaussWeights[i] * ((2 * Math.PI) / nphi);
135 }
136 // twiddle/phase table in f64 (device sin/cos is too inaccurate: ~2^-11 under Vulkan)
137 const trig = new Float32Array(2 * nphi);
138 for (let k = 0; k < nphi; k++) {
139 trig[2 * k] = Math.cos((2 * Math.PI * k) / nphi);
140 trig[2 * k + 1] = Math.sin((2 * Math.PI * k) / nphi);
141 }
143 const mkBuf = (label: string, size: number, usage: GPUBufferUsageFlags) =>
144 dev.createBuffer({ label, size, usage });
145 this.bufAb = mkBuf('sht-ab', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
146 this.bufAmm = mkBuf('sht-amm', 4 * (mmax + 1), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
147 this.bufCtstw = mkBuf('sht-ctstw', 4 * 3 * nlat, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
148 this.bufTrig = mkBuf('sht-trig', 8 * nphi, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
149 self.qlmIn = mkBuf('sht-qlm-in', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC);
150 self.qlmOut = mkBuf('sht-qlm-out', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
151 self.fmBuf = mkBuf('sht-fm', 8 * (mmax + 1) * nlat, GPUBufferUsage.STORAGE);
152 self.spatBuf = mkBuf('sht-spat', 4 * nlat * nphi, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC);
153 this.stageSpat = mkBuf('sht-stage-spat', 4 * nlat * nphi, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
154 this.stageQ = mkBuf('sht-stage-q', 8 * this.nlm, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
156 dev.queue.writeBuffer(this.bufAb, 0, new Float32Array(ab));
157 dev.queue.writeBuffer(this.bufAmm, 0, new Float32Array(amm));
158 dev.queue.writeBuffer(this.bufCtstw, 0, ctstw);
159 dev.queue.writeBuffer(this.bufTrig, 0, trig);
161 // --- shaders / pipelines ---
162 const legP = { lmax, mmax, nlat, wgSynth: WG_SYNTH, wgAnalys: WG_ANALYS };
163 const fourP = { mmax, nlat, nphi };
164 const [pLegS, pLegA, pFourS, pFourA] = await Promise.all([
165 makePipeline(dev, legSynthWGSL(legP), 'leg_synth'),
166 makePipeline(dev, legAnalysWGSL(legP), 'leg_analys'),
167 makePipeline(
168 dev,
169 this.fourierMode === 'fft' ? fftSynthWGSL(fourP) : dftSynthWGSL(fourP),
170 this.fourierMode === 'fft' ? 'fft_synth' : 'dft_synth',
171 ),
172 makePipeline(
173 dev,
174 this.fourierMode === 'fft' ? fftAnalysWGSL(fourP) : dftAnalysWGSL(fourP),
175 this.fourierMode === 'fft' ? 'fft_analys' : 'dft_analys',
176 ),
177 ]);
178 this.pipeLegSynth = pLegS;
179 this.pipeLegAnalys = pLegA;
180 this.pipeFourSynth = pFourS;
181 this.pipeFourAnalys = pFourA;
183 const entries = (bufs: GPUBuffer[]) =>
184 bufs.map((buffer, binding) => ({ binding, resource: { buffer } }));
185 this.bgLegSynth = dev.createBindGroup({
186 layout: pLegS.getBindGroupLayout(0),
187 entries: entries([this.bufAb, this.bufAmm, this.bufCtstw, this.qlmIn, this.fmBuf]),
188 });
189 this.bgLegAnalys = dev.createBindGroup({
190 layout: pLegA.getBindGroupLayout(0),
191 entries: entries([this.bufAb, this.bufAmm, this.bufCtstw, this.fmBuf, this.qlmOut]),
192 });
193 this.bgFourSynth = dev.createBindGroup({
194 layout: pFourS.getBindGroupLayout(0),
195 entries: entries([this.fmBuf, this.spatBuf, this.bufTrig]),
196 });
197 this.bgFourAnalys = dev.createBindGroup({
198 layout: pFourA.getBindGroupLayout(0),
199 entries: entries([this.spatBuf, this.fmBuf, this.bufTrig]),
200 });
201 }
203 /** Record the synthesis (spectral qlmIn -> spatial spatBuf) into an encoder. */
204 encodeSynth(encoder: GPUCommandEncoder): void {
205 const { mmax, nlat, nphi } = this.cfg;
206 const pass = encoder.beginComputePass({ label: 'sht-synth' });
207 pass.setPipeline(this.pipeLegSynth);
208 pass.setBindGroup(0, this.bgLegSynth);
209 pass.dispatchWorkgroups(Math.ceil(nlat / WG_SYNTH), mmax + 1);
210 pass.setPipeline(this.pipeFourSynth);
211 pass.setBindGroup(0, this.bgFourSynth);
212 if (this.fourierMode === 'fft') {
213 pass.dispatchWorkgroups(nlat);
214 } else {
215 pass.dispatchWorkgroups(Math.ceil(nphi / 64), nlat);
216 }
217 pass.end();
218 }
220 /** Record the analysis (spatial spatBuf -> spectral qlmOut) into an encoder. */
221 encodeAnalys(encoder: GPUCommandEncoder): void {
222 const { mmax, nlat } = this.cfg;
223 const pass = encoder.beginComputePass({ label: 'sht-analys' });
224 pass.setPipeline(this.pipeFourAnalys);
225 pass.setBindGroup(0, this.bgFourAnalys);
226 if (this.fourierMode === 'fft') {
227 pass.dispatchWorkgroups(nlat);
228 } else {
229 pass.dispatchWorkgroups(Math.ceil((mmax + 1) / 64), nlat);
230 }
231 pass.setPipeline(this.pipeLegAnalys);
232 pass.setBindGroup(0, this.bgLegAnalys);
233 pass.dispatchWorkgroups(mmax + 1);
234 pass.end();
235 }
237 /**
238 * Spectral -> spatial. qlm: interleaved [re, im], SHTNS LM ordering,
239 * length 2*nlm. Returns the spatial field, length nlat*nphi.
240 */
241 async synth(qlm: Float32Array): Promise<Float32Array> {
242 const { nlat, nphi } = this.cfg;
243 if (qlm.length !== 2 * this.nlm) throw new Error(`qlm must have length ${2 * this.nlm}`);
244 this.device.queue.writeBuffer(this.qlmIn, 0, qlm as Float32Array<ArrayBuffer>);
245 const enc = this.device.createCommandEncoder();
246 this.encodeSynth(enc);
247 enc.copyBufferToBuffer(this.spatBuf, 0, this.stageSpat, 0, 4 * nlat * nphi);
248 this.device.queue.submit([enc.finish()]);
249 await this.stageSpat.mapAsync(GPUMapMode.READ);
250 const out = new Float32Array(this.stageSpat.getMappedRange().slice(0));
251 this.stageSpat.unmap();
252 return out;
253 }
255 /** Spatial -> spectral. spat: length nlat*nphi. Returns interleaved qlm, length 2*nlm. */
256 async analys(spat: Float32Array): Promise<Float32Array> {
257 const { nlat, nphi } = this.cfg;
258 if (spat.length !== nlat * nphi) throw new Error(`spat must have length ${nlat * nphi}`);
259 this.device.queue.writeBuffer(this.spatBuf, 0, spat as Float32Array<ArrayBuffer>);
260 const enc = this.device.createCommandEncoder();
261 this.encodeAnalys(enc);
262 enc.copyBufferToBuffer(this.qlmOut, 0, this.stageQ, 0, 8 * this.nlm);
263 this.device.queue.submit([enc.finish()]);
264 await this.stageQ.mapAsync(GPUMapMode.READ);
265 const out = new Float32Array(this.stageQ.getMappedRange().slice(0));
266 this.stageQ.unmap();
267 return out;
268 }
270 destroy(): void {
271 for (const b of [
272 this.bufAb, this.bufAmm, this.bufCtstw, this.bufTrig, this.qlmIn, this.qlmOut,
273 this.fmBuf, this.spatBuf, this.stageSpat, this.stageQ,
274 ]) b?.destroy();
275 }
276}
278/** Request an adapter/device suitable for the transforms. */
279export async function requestShtDevice(): Promise<GPUDevice> {
280 if (!navigator.gpu) throw new Error('WebGPU is not available in this browser');
281 const adapter = await navigator.gpu.requestAdapter();
282 if (!adapter) throw new Error('No WebGPU adapter available');
283 // ask for a larger workgroup storage if the adapter offers it (bigger FFTs)
284 const wgStorage = Math.min(adapter.limits.maxComputeWorkgroupStorageSize, 32768);
285 return adapter.requestDevice({
286 requiredLimits: { maxComputeWorkgroupStorageSize: wgStorage },
287 });
288}