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';
13 legSynthWGSL,
14 legAnalysWGSL,
15 legSynthBatchWGSL,
16 legAnalysBatchWGSL,
17} from './wgsl/leg.ts';
0d99c91Differentiate the phi flux in grid spaceDan Fortunato 18import { fmDphiWGSL } from './wgsl/deriv.ts';
20 fftSynthWGSL,
21 fftAnalysWGSL,
22 fftSynthRealWGSL,
23 fftAnalysRealWGSL,
24 dftSynthWGSL,
25 dftAnalysWGSL,
26 fftThreads,
27} from './wgsl/fourier.ts';
29export type FourierMode = 'auto' | 'fft' | 'dft';
31/** The two bind groups (Legendre stage, Fourier stage) of one transform. */
32export interface ShtBinding {
33 readonly bgLeg: GPUBindGroup;
34 readonly bgFour: GPUBindGroup;
35}
0d99c91Differentiate the phi flux in grid spaceDan Fortunato 37/** The three bind groups of one grid-space phi-derivative (see dphig). */
38export interface ShtDphigBinding {
39 readonly bgFourAnalys: GPUBindGroup;
40 readonly bgMul: GPUBindGroup;
41 readonly bgFourSynth: GPUBindGroup;
42}
45 * One batched transform: K fields through a single Legendre dispatch (the
46 * recurrence walked once, K accumulator lanes) plus K per-field Fourier
47 * dispatches — the Fourier stage shares nothing across fields, so batching
48 * it would save only bind-group switches.
49 */
50export interface ShtBatchBinding {
51 /** Lanes in this batch — selects the pipeline compiled for that width. */
52 readonly size: number;
53 readonly bgLeg: GPUBindGroup;
54 /** Per-lane Fourier bind group, lane k against fm arena k. */
55 readonly bgFour: GPUBindGroup[];
56}
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 58const bgEntries = (bufs: GPUBuffer[]) =>
59 bufs.map((buffer, binding) => ({ binding, resource: { buffer } }));
61export interface ShtOptions {
62 /** Fourier stage implementation. 'auto' picks fft when nphi is a power of two that fits in workgroup memory. */
63 fourier?: FourierMode;
64}
66const WG_SYNTH = 64;
68/**
69 * Tuning knob, for A/B-ing a change without editing code. Reads globalThis
70 * first (set it before creating a plan, as scripts/_ab.ts does), then the
71 * environment, so `SHT_SUBGROUPS=0 npm run bench:sht` works too. `process` is
72 * absent in the browser, where only the globalThis form applies.
73 */
74function tuning(name: string): unknown {
75 const g = (globalThis as Record<string, unknown>)[name];
76 if (g !== undefined) return g;
77 const env = (globalThis as { process?: { env?: Record<string, string | undefined> } }).process?.env?.[name];
78 if (env === undefined || env === '') return undefined;
79 if (env === '1' || env === 'true') return true;
80 if (env === '0' || env === 'false') return false;
81 const n = Number(env);
82 return Number.isFinite(n) ? n : env;
83}
85/**
86 * Workgroup size for the analysis Legendre reduction. The right answer differs
87 * between the two reduction strategies, so it is chosen per strategy.
88 *
89 * Measured on an RTX PRO 6000 Blackwell (analysis, us). Shared-memory tree,
90 * where each doubling of wgAnalys costs another barrier per l-pair:
91 *
92 * wgAnalys: 16 32 64 128 256
93 * nlat=128 43.2 40.6 42.6 46.7 52.9 -> 32
94 * nlat=256 104.0 85.6 87.8 89.9 100.0 -> 32
95 * nlat=512 408.8 216.9 184.4 190.8 204.8 -> 64
96 *
97 * i.e. max(32, nlat/8). A flat 32 would be worse than the old default of 256 at
98 * nlat=512, so it cannot be fitted on one grid. With subgroupAdd the barrier
99 * count stops growing with wgAnalys and the picture inverts: threads in flight,
100 * (mmax+1) * wgAnalys, becomes binding, since analysis dispatches only mmax+1
101 * workgroups. 128 then wins at every grid (round trip, us):
102 *
103 * 128x256 37.8 (vs 38.3), 256x512 66.6 (vs 74.7), 512x1024 131.9 (vs 156.6)
104 */
105function defaultWgAnalys(nlat: number, limit: number, subgroups: boolean): number {
106 if (subgroups) {
107 // capped at nlat so small grids do not launch threads with no latitude to own
108 let cap = 1;
109 while (cap < nlat) cap *= 2;
110 return Math.min(128, limit, cap);
111 }
112 const target = Math.max(32, nlat / 8);
113 let wg = 1;
114 while (wg < target) wg *= 2; // the tree reduction halves, so a power of two
115 return Math.min(wg, limit);
116}
118async function makePipeline(
119 device: GPUDevice,
120 code: string,
121 entryPoint: string,
122): Promise<GPUComputePipeline> {
123 device.pushErrorScope('validation');
124 const module = device.createShaderModule({ code, label: entryPoint });
125 const info = await module.getCompilationInfo();
126 const errors = info.messages.filter((m) => m.type === 'error');
127 if (errors.length) {
128 throw new Error(
129 `WGSL compile error in ${entryPoint}:\n` +
130 errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n'),
131 );
132 }
133 const pipeline = await device.createComputePipelineAsync({
134 layout: 'auto',
135 compute: { module, entryPoint },
136 label: entryPoint,
137 });
138 const err = await device.popErrorScope();
139 if (err) throw new Error(`pipeline ${entryPoint}: ${err.message}`);
140 return pipeline;
141}
143export class ShtPlan {
144 readonly cfg: ShtConfig;
145 readonly nlm: number;
146 readonly fourierMode: 'fft' | 'dft';
147 /** Latitudes leg_synth walks: nlat/2 when parity folding. */
148 readonly legLat: number = 0;
150 * Widest transform batch this plan supports: the largest even K <= 4 whose
151 * Legendre bind group (3 tables + K caller fields + the shared fm arena)
152 * fits the device's storage-buffer limit. K = 4 needs exactly the WebGPU
153 * default of 8, so batching is fully available on every stack; 1 (no
154 * batching) if SHT_BATCH is turned off. Batched and scalar transforms
155 * compute identical per-lane arithmetic, so this only affects speed,
156 * never results.
157 */
158 readonly batchK: number = 1;
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 159 /** Colatitudes theta_i (f64, increasing: north to south). */
160 readonly theta: Float64Array;
161 readonly cosTheta: Float64Array;
162 readonly gaussWeights: Float64Array;
164 private device: GPUDevice;
165 private bufAb!: GPUBuffer;
166 private bufAmm!: GPUBuffer;
167 private bufCtstw!: GPUBuffer;
168 private bufTrig!: GPUBuffer;
169 /** Spectral input (synthesis) — write with queue.writeBuffer or use synth(). */
170 readonly qlmIn!: GPUBuffer;
171 /** Spectral output (analysis). */
172 readonly qlmOut!: GPUBuffer;
173 /** Fourier-space intermediate [(m)*nlat + ilat], complex f32. COPY_SRC so the
174 * stage boundary is observable: a transform is Legendre-then-Fourier, and
175 * scripts/diagnose-sht.ts tells the two apart by reading this. */
176 readonly fmBuf!: GPUBuffer;
177 /** Spatial field [ilat*nphi + iphi], f32. */
178 readonly spatBuf!: GPUBuffer;
179 private stageSpat!: GPUBuffer;
180 private stageQ!: GPUBuffer;
182 private pipeLegSynth!: GPUComputePipeline;
183 private pipeLegAnalys!: GPUComputePipeline;
184 private pipeFourSynth!: GPUComputePipeline;
185 private pipeFourAnalys!: GPUComputePipeline;
0d99c91Differentiate the phi flux in grid spaceDan Fortunato 186 /** Fourier-space i*m multiply, the middle of dphig. */
187 private pipeFmDphi!: GPUComputePipeline;
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 188 /** Batched Legendre pipelines by lane count (even sizes up to batchK). */
189 private pipeLegSynthB = new Map<number, GPUComputePipeline>();
190 private pipeLegAnalysB = new Map<number, GPUComputePipeline>();
191 /** One fm arena for all batch lanes (lane k at byte offset k * fmLaneBytes,
192 * 256-aligned so the Fourier stage can bind a lane by buffer offset). A
193 * single buffer keeps the batched Legendre bind group at 3 tables +
194 * K fields + 1 arena — within WebGPU's default storage-buffer limit of 8
195 * at K = 4, on every stack. */
196 private fmArena: GPUBuffer | null = null;
197 private fmLaneBytes = 0;
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 198 private bgLegSynth!: GPUBindGroup;
199 private bgLegAnalys!: GPUBindGroup;
200 private bgFourSynth!: GPUBindGroup;
201 private bgFourAnalys!: GPUBindGroup;
203 private constructor(device: GPUDevice, cfg: ShtConfig, fourierMode: 'fft' | 'dft') {
204 this.device = device;
205 this.cfg = cfg;
206 this.nlm = nlmCalc(cfg.lmax, cfg.mmax);
207 this.fourierMode = fourierMode;
208 const { x, w } = gaussNodesWeights(cfg.nlat);
209 this.cosTheta = x;
210 this.gaussWeights = w;
211 this.theta = new Float64Array(cfg.nlat);
212 for (let i = 0; i < cfg.nlat; i++) this.theta[i] = Math.acos(x[i]);
213 }
215 static async create(device: GPUDevice, cfg: ShtConfig, opts: ShtOptions = {}): Promise<ShtPlan> {
216 validateConfig(cfg);
217 const want = opts.fourier ?? 'auto';
218 const fftFits =
219 isPowerOfTwo(cfg.nphi) &&
220 16 * cfg.nphi <= device.limits.maxComputeWorkgroupStorageSize &&
221 fftThreads(cfg.nphi) <= device.limits.maxComputeInvocationsPerWorkgroup;
222 if (want === 'fft' && !fftFits) {
223 throw new Error(
224 `fourier:'fft' requires power-of-two nphi with 16*nphi <= maxComputeWorkgroupStorageSize ` +
225 `(nphi=${cfg.nphi}, limit=${device.limits.maxComputeWorkgroupStorageSize})`,
226 );
227 }
228 const mode: 'fft' | 'dft' = want === 'dft' ? 'dft' : fftFits ? 'fft' : 'dft';
229 const plan = new ShtPlan(device, cfg, mode);
230 await plan.init();
231 return plan;
232 }
234 private async init(): Promise<void> {
235 const { lmax, mmax, nlat, nphi } = this.cfg;
236 const dev = this.device;
237 const self = this as {
238 -readonly [k in keyof ShtPlan]: ShtPlan[k];
239 };
241 // --- host precomputation (f64), then downcast to f32 for upload ---
242 const { amm, ab } = legendreCoeffs(lmax, mmax);
243 const ctstw = new Float32Array(3 * nlat);
244 for (let i = 0; i < nlat; i++) {
245 ctstw[i] = this.cosTheta[i];
246 ctstw[nlat + i] = Math.sqrt(1 - this.cosTheta[i] * this.cosTheta[i]);
247 ctstw[2 * nlat + i] = this.gaussWeights[i] * ((2 * Math.PI) / nphi);
248 }
249 // twiddle/phase table in f64 (device sin/cos is too inaccurate: ~2^-11 under Vulkan)
250 const trig = new Float32Array(2 * nphi);
251 for (let k = 0; k < nphi; k++) {
252 trig[2 * k] = Math.cos((2 * Math.PI * k) / nphi);
253 trig[2 * k + 1] = Math.sin((2 * Math.PI * k) / nphi);
254 }
256 const mkBuf = (label: string, size: number, usage: GPUBufferUsageFlags) =>
257 dev.createBuffer({ label, size, usage });
258 this.bufAb = mkBuf('sht-ab', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
259 this.bufAmm = mkBuf('sht-amm', 4 * (mmax + 1), GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
260 this.bufCtstw = mkBuf('sht-ctstw', 4 * 3 * nlat, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
261 this.bufTrig = mkBuf('sht-trig', 8 * nphi, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
262 self.qlmIn = mkBuf('sht-qlm-in', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC);
263 self.qlmOut = mkBuf('sht-qlm-out', 8 * this.nlm, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
264 self.fmBuf = mkBuf('sht-fm', 8 * (mmax + 1) * nlat, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
265 self.spatBuf = mkBuf('sht-spat', 4 * nlat * nphi, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC);
266 this.stageSpat = mkBuf('sht-stage-spat', 4 * nlat * nphi, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
267 this.stageQ = mkBuf('sht-stage-q', 8 * this.nlm, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
269 dev.queue.writeBuffer(this.bufAb, 0, new Float32Array(ab));
270 dev.queue.writeBuffer(this.bufAmm, 0, new Float32Array(amm));
271 dev.queue.writeBuffer(this.bufCtstw, 0, ctstw);
272 dev.queue.writeBuffer(this.bufTrig, 0, trig);
274 // --- shaders / pipelines ---
275 const subgroups = tuning('SHT_SUBGROUPS') !== false && dev.features.has('subgroups');
276 // parity folding needs an equator-symmetric grid; Gauss nodes are, if nlat is even
277 const parity = tuning('SHT_PARITY') !== false && nlat % 2 === 0;
278 const wgAnalys =
279 (tuning('SHT_WG_ANALYS') as number | undefined) ??
280 defaultWgAnalys(nlat, dev.limits.maxComputeInvocationsPerWorkgroup, subgroups);
281 const legP = {
282 lmax,
283 mmax,
284 nlat,
285 wgSynth: WG_SYNTH,
286 wgAnalys,
287 subgroups,
288 spanPairs: tuning('SHT_SPAN_PAIRS') as number | undefined,
289 parity,
290 };
291 (this as { legLat: number }).legLat = parity ? nlat / 2 : nlat;
292 const fourP = { mmax, nlat, nphi, radix: (tuning('SHT_RADIX') as number | undefined) ?? 4 };
293 // The spatial field is real (layout.ts stores m >= 0 only), so the Fourier
294 // stage can run an nphi/2-point complex FFT plus a recombination instead of
295 // a full nphi-point one: half the arithmetic and half the workgroup storage.
296 // The complex kernels remain for a future complex-valued field, and are what
297 // SHT_REAL_FFT=0 selects.
298 const realFft =
299 this.fourierMode === 'fft' && nphi % 2 === 0 && tuning('SHT_REAL_FFT') !== false;
300 const fftS = realFft ? fftSynthRealWGSL : fftSynthWGSL;
301 const fftA = realFft ? fftAnalysRealWGSL : fftAnalysWGSL;
0d99c91Differentiate the phi flux in grid spaceDan Fortunato 302 const [pLegS, pLegA, pFourS, pFourA, pFmDphi] = await Promise.all([
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 303 makePipeline(dev, legSynthWGSL(legP), 'leg_synth'),
304 makePipeline(dev, legAnalysWGSL(legP), 'leg_analys'),
305 makePipeline(
306 dev,
307 this.fourierMode === 'fft' ? fftS(fourP) : dftSynthWGSL(fourP),
308 this.fourierMode === 'fft' ? 'fft_synth' : 'dft_synth',
309 ),
310 makePipeline(
311 dev,
312 this.fourierMode === 'fft' ? fftA(fourP) : dftAnalysWGSL(fourP),
313 this.fourierMode === 'fft' ? 'fft_analys' : 'dft_analys',
314 ),
0d99c91Differentiate the phi flux in grid spaceDan Fortunato 315 // Mirrors filterMask: content at l >= lmax-2 is filtered on the
316 // l-space route, so the m-space route keeps m <= lmax-3.
317 makePipeline(
318 dev,
319 fmDphiWGSL({ mmax, nlat, nphi, mcut: lmax - 3 }),
320 'fm_dphi',
321 ),
323 this.pipeLegSynth = pLegS;
324 this.pipeLegAnalys = pLegA;
325 this.pipeFourSynth = pFourS;
326 this.pipeFourAnalys = pFourA;
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 328
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 329 // --- batched Legendre pipelines ---
330 // The widest even K <= 4 whose bind group (3 tables + K fields + the fm
331 // arena) fits the device's storage-buffer budget — K = 4 needs 8, the
332 // WebGPU default, so batching is fully available everywhere unless
333 // SHT_BATCH=0 disables it (SHT_BATCH=2 caps it, for A/B).
334 const batchTuning = tuning('SHT_BATCH');
335 const batchWant =
336 batchTuning === false || batchTuning === 0
337 ? 1
338 : typeof batchTuning === 'number'
339 ? batchTuning
340 : 4;
341 const batchFit = dev.limits.maxStorageBuffersPerShaderStage - 4;
342 const batchK = Math.min(4, Math.max(1, batchWant), 2 * Math.floor(batchFit / 2));
343 (this as { batchK: number }).batchK = batchK;
344 if (batchK >= 2) {
345 // Lane stride rounded to the 256-byte offset alignment buffer bindings
346 // require; laneElems is that stride in vec2f units for the kernels.
347 this.fmLaneBytes = Math.ceil((8 * (mmax + 1) * nlat) / 256) * 256;
348 const laneElems = this.fmLaneBytes / 8;
349 this.fmArena = mkBuf('sht-fm-arena', batchK * this.fmLaneBytes, GPUBufferUsage.STORAGE);
350 const sizes = [];
351 for (let k = 2; k <= batchK; k += 2) sizes.push(k);
352 const pipes = await Promise.all(
353 sizes.flatMap((k) => [
354 makePipeline(dev, legSynthBatchWGSL(legP, k, laneElems), `leg_synth_batch`),
355 makePipeline(dev, legAnalysBatchWGSL(legP, k, laneElems), `leg_analys_batch`),
356 ]),
357 );
358 sizes.forEach((k, i) => {
359 this.pipeLegSynthB.set(k, pipes[2 * i]);
360 this.pipeLegAnalysB.set(k, pipes[2 * i + 1]);
361 });
362 }
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 364 const entries = bgEntries;
365 this.bgLegSynth = dev.createBindGroup({
366 layout: pLegS.getBindGroupLayout(0),
367 entries: entries([this.bufAb, this.bufAmm, this.bufCtstw, this.qlmIn, this.fmBuf]),
368 });
369 this.bgLegAnalys = dev.createBindGroup({
370 layout: pLegA.getBindGroupLayout(0),
371 entries: entries([this.bufAb, this.bufAmm, this.bufCtstw, this.fmBuf, this.qlmOut]),
372 });
373 this.bgFourSynth = dev.createBindGroup({
374 layout: pFourS.getBindGroupLayout(0),
375 entries: entries([this.fmBuf, this.spatBuf, this.bufTrig]),
376 });
377 this.bgFourAnalys = dev.createBindGroup({
378 layout: pFourA.getBindGroupLayout(0),
379 entries: entries([this.spatBuf, this.fmBuf, this.bufTrig]),
380 });
381 }
383 /**
384 * Bind groups for one transform against caller-supplied spectral/spatial
385 * buffers, so a transform can read and write buffers it does not own (the
386 * .m-driven executor keeps a buffer per IR variable). Build these once at
387 * plan time, not per step. `fmBuf` stays internal scratch: passes and
388 * dispatches within a submission execute in order, so sequential transforms
389 * can share it.
390 */
391 createSynthBinding(qlmIn: GPUBuffer, spatOut: GPUBuffer): ShtBinding {
392 return {
393 bgLeg: this.device.createBindGroup({
394 layout: this.pipeLegSynth.getBindGroupLayout(0),
395 entries: bgEntries([this.bufAb, this.bufAmm, this.bufCtstw, qlmIn, this.fmBuf]),
396 }),
397 bgFour: this.device.createBindGroup({
398 layout: this.pipeFourSynth.getBindGroupLayout(0),
399 entries: bgEntries([this.fmBuf, spatOut, this.bufTrig]),
400 }),
401 };
402 }
404 createAnalysBinding(spatIn: GPUBuffer, qlmOut: GPUBuffer): ShtBinding {
405 return {
406 bgFour: this.device.createBindGroup({
407 layout: this.pipeFourAnalys.getBindGroupLayout(0),
408 entries: bgEntries([spatIn, this.fmBuf, this.bufTrig]),
409 }),
410 bgLeg: this.device.createBindGroup({
411 layout: this.pipeLegAnalys.getBindGroupLayout(0),
412 entries: bgEntries([this.bufAb, this.bufAmm, this.bufCtstw, this.fmBuf, qlmOut]),
413 }),
414 };
415 }
418 * Bind groups for one batched synthesis: members.length must be a compiled
419 * lane count (an even size <= batchK). Member outputs must be distinct
420 * buffers; each lane gets its own fm arena, so batches compose in a pass
421 * exactly like sequential scalar transforms do.
422 */
423 /** The fm arena sliced at lane k, sized as one transform's fm. */
424 #fmLane(k: number): GPUBufferBinding {
425 const { mmax, nlat } = this.cfg;
426 return {
427 buffer: this.fmArena!,
428 offset: k * this.fmLaneBytes,
429 size: 8 * (mmax + 1) * nlat,
430 };
431 }
433 createSynthBatchBinding(
434 members: { qlmIn: GPUBuffer; spatOut: GPUBuffer }[],
435 ): ShtBatchBinding {
436 const K = members.length;
437 const pipe = this.pipeLegSynthB.get(K);
438 if (!pipe) throw new Error(`no batched synthesis pipeline for ${K} lanes`);
439 return {
440 size: K,
441 bgLeg: this.device.createBindGroup({
442 layout: pipe.getBindGroupLayout(0),
443 entries: [
444 ...bgEntries([this.bufAb, this.bufAmm, this.bufCtstw, ...members.map((m) => m.qlmIn)]),
445 { binding: 3 + K, resource: { buffer: this.fmArena! } },
446 ],
447 }),
448 bgFour: members.map((m, k) =>
449 this.device.createBindGroup({
450 layout: this.pipeFourSynth.getBindGroupLayout(0),
451 entries: [
452 { binding: 0, resource: this.#fmLane(k) },
453 { binding: 1, resource: { buffer: m.spatOut } },
454 { binding: 2, resource: { buffer: this.bufTrig } },
455 ],
456 }),
457 ),
458 };
459 }
461 createAnalysBatchBinding(
462 members: { spatIn: GPUBuffer; qlmOut: GPUBuffer }[],
463 ): ShtBatchBinding {
464 const K = members.length;
465 const pipe = this.pipeLegAnalysB.get(K);
466 if (!pipe) throw new Error(`no batched analysis pipeline for ${K} lanes`);
467 return {
468 size: K,
469 bgFour: members.map((m, k) =>
470 this.device.createBindGroup({
471 layout: this.pipeFourAnalys.getBindGroupLayout(0),
472 entries: [
473 { binding: 0, resource: { buffer: m.spatIn } },
474 { binding: 1, resource: this.#fmLane(k) },
475 { binding: 2, resource: { buffer: this.bufTrig } },
476 ],
477 }),
478 ),
479 bgLeg: this.device.createBindGroup({
480 layout: pipe.getBindGroupLayout(0),
481 entries: [
482 ...bgEntries([this.bufAb, this.bufAmm, this.bufCtstw]),
483 { binding: 3, resource: { buffer: this.fmArena! } },
484 ...members.map((m, k) => ({
485 binding: 4 + k,
486 resource: { buffer: m.qlmOut },
487 })),
488 ],
489 }),
490 };
491 }
494 * Bind groups for one grid-space phi-derivative, dphig: Fourier analysis
495 * of each latitude row into fm (which truncates to m <= mmax for free),
496 * the i*m/NPHI multiply (zeroing m past the top-degree filt's reach), and
497 * Fourier synthesis back to the grid. No Legendre stage anywhere — this
498 * is what lets the flux-form divergence drop the Q-flux's spherical-
499 * harmonic analysis (docs/reduced-transforms.md Sec 5b's companion trick
500 * in Sec 6-of-changes): d/dphi is diagonal in the Fourier index. Uses
501 * fmBuf as scratch, sequentially like every transform in a pass.
502 */
503 createDphigBinding(spatIn: GPUBuffer, spatOut: GPUBuffer): ShtDphigBinding {
504 return {
505 bgFourAnalys: this.device.createBindGroup({
506 layout: this.pipeFourAnalys.getBindGroupLayout(0),
507 entries: bgEntries([spatIn, this.fmBuf, this.bufTrig]),
508 }),
509 bgMul: this.device.createBindGroup({
510 layout: this.pipeFmDphi.getBindGroupLayout(0),
511 entries: bgEntries([this.fmBuf]),
512 }),
513 bgFourSynth: this.device.createBindGroup({
514 layout: this.pipeFourSynth.getBindGroupLayout(0),
515 entries: bgEntries([this.fmBuf, spatOut, this.bufTrig]),
516 }),
517 };
518 }
520 /** Record dphig into an existing compute pass: two Fourier stages and a
521 * pointwise multiply — no Legendre work. */
522 encodeDphigInto(pass: GPUComputePassEncoder, b: ShtDphigBinding): void {
523 const { mmax, nlat, nphi } = this.cfg;
524 pass.setPipeline(this.pipeFourAnalys);
525 pass.setBindGroup(0, b.bgFourAnalys);
526 if (this.fourierMode === 'fft') {
527 pass.dispatchWorkgroups(nlat);
528 } else {
529 pass.dispatchWorkgroups(Math.ceil((mmax + 1) / 64), nlat);
530 }
531 pass.setPipeline(this.pipeFmDphi);
532 pass.setBindGroup(0, b.bgMul);
533 pass.dispatchWorkgroups(Math.ceil(((mmax + 1) * nlat) / 64));
534 pass.setPipeline(this.pipeFourSynth);
535 pass.setBindGroup(0, b.bgFourSynth);
536 if (this.fourierMode === 'fft') {
537 pass.dispatchWorkgroups(nlat);
538 } else {
539 pass.dispatchWorkgroups(Math.ceil(nphi / 64), nlat);
540 }
541 }
543 /** CPU convenience: grid field -> d/dphi of its trig interpolant, for tests. */
544 async dphig(spat: Float32Array): Promise<Float32Array> {
545 const { nlat, nphi } = this.cfg;
546 if (spat.length !== nlat * nphi) throw new Error(`spat must have length ${nlat * nphi}`);
547 this.device.queue.writeBuffer(this.spatBuf, 0, spat as Float32Array<ArrayBuffer>);
548 const binding = this.createDphigBinding(this.spatBuf, this.spatBuf);
549 const enc = this.device.createCommandEncoder({ label: 'sht-dphig' });
550 const pass = enc.beginComputePass({ label: 'sht-dphig' });
551 this.encodeDphigInto(pass, binding);
552 pass.end();
553 enc.copyBufferToBuffer(this.spatBuf, 0, this.stageSpat, 0, 4 * nlat * nphi);
554 this.device.queue.submit([enc.finish()]);
555 await this.stageSpat.mapAsync(GPUMapMode.READ);
556 const out = new Float32Array(this.stageSpat.getMappedRange().slice(0));
557 this.stageSpat.unmap();
558 return out;
559 }
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 561 /** Record a batched synthesis: one Legendre dispatch, K Fourier dispatches. */
562 encodeSynthBatchInto(pass: GPUComputePassEncoder, b: ShtBatchBinding): void {
563 const { mmax, nlat, nphi } = this.cfg;
564 pass.setPipeline(this.pipeLegSynthB.get(b.size)!);
565 pass.setBindGroup(0, b.bgLeg);
566 pass.dispatchWorkgroups(Math.ceil(this.legLat / WG_SYNTH), mmax + 1);
567 pass.setPipeline(this.pipeFourSynth);
568 for (const bg of b.bgFour) {
569 pass.setBindGroup(0, bg);
570 if (this.fourierMode === 'fft') {
571 pass.dispatchWorkgroups(nlat);
572 } else {
573 pass.dispatchWorkgroups(Math.ceil(nphi / 64), nlat);
574 }
575 }
576 }
578 /** Record a batched analysis: K Fourier dispatches, one Legendre dispatch. */
579 encodeAnalysBatchInto(pass: GPUComputePassEncoder, b: ShtBatchBinding): void {
580 const { mmax, nlat } = this.cfg;
581 pass.setPipeline(this.pipeFourAnalys);
582 for (const bg of b.bgFour) {
583 pass.setBindGroup(0, bg);
584 if (this.fourierMode === 'fft') {
585 pass.dispatchWorkgroups(nlat);
586 } else {
587 pass.dispatchWorkgroups(Math.ceil((mmax + 1) / 64), nlat);
588 }
589 }
590 pass.setPipeline(this.pipeLegAnalysB.get(b.size)!);
591 pass.setBindGroup(0, b.bgLeg);
592 pass.dispatchWorkgroups(mmax + 1);
593 }
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 595 /** Record synthesis into an existing compute pass. */
596 encodeSynthInto(pass: GPUComputePassEncoder, b: ShtBinding): void {
597 const { mmax, nlat, nphi } = this.cfg;
598 pass.setPipeline(this.pipeLegSynth);
599 pass.setBindGroup(0, b.bgLeg);
600 pass.dispatchWorkgroups(Math.ceil(this.legLat / WG_SYNTH), mmax + 1);
601 pass.setPipeline(this.pipeFourSynth);
602 pass.setBindGroup(0, b.bgFour);
603 if (this.fourierMode === 'fft') {
604 pass.dispatchWorkgroups(nlat);
605 } else {
606 pass.dispatchWorkgroups(Math.ceil(nphi / 64), nlat);
607 }
608 }
610 /** Record analysis into an existing compute pass. */
611 encodeAnalysInto(pass: GPUComputePassEncoder, b: ShtBinding): void {
612 const { mmax, nlat } = this.cfg;
613 pass.setPipeline(this.pipeFourAnalys);
614 pass.setBindGroup(0, b.bgFour);
615 if (this.fourierMode === 'fft') {
616 pass.dispatchWorkgroups(nlat);
617 } else {
618 pass.dispatchWorkgroups(Math.ceil((mmax + 1) / 64), nlat);
619 }
620 pass.setPipeline(this.pipeLegAnalys);
621 pass.setBindGroup(0, b.bgLeg);
622 pass.dispatchWorkgroups(mmax + 1);
623 }
625 /** Record the synthesis (spectral qlmIn -> spatial spatBuf) into an encoder. */
626 encodeSynth(encoder: GPUCommandEncoder): void {
627 const pass = encoder.beginComputePass({ label: 'sht-synth' });
628 this.encodeSynthInto(pass, { bgLeg: this.bgLegSynth, bgFour: this.bgFourSynth });
629 pass.end();
630 }
632 /**
633 * Diagnostics: encode one stage alone, in its own pass, so a timestamp query
634 * can measure just that kernel. The solver wants both stages in a shared pass
635 * and should use encodeSynthInto/encodeAnalysInto; this exists because
636 * inferring per-kernel cost by subtracting trivially-sized runs is unreliable.
637 */
638 encodeStage(
639 encoder: GPUCommandEncoder,
640 stage: 'legSynth' | 'fourSynth' | 'fourAnalys' | 'legAnalys',
641 timestampWrites?: GPUComputePassTimestampWrites,
642 ): void {
643 const { mmax, nlat, nphi } = this.cfg;
644 const fft = this.fourierMode === 'fft';
645 const pass = encoder.beginComputePass({ label: `sht-${stage}`, timestampWrites });
646 switch (stage) {
647 case 'legSynth':
648 pass.setPipeline(this.pipeLegSynth);
649 pass.setBindGroup(0, this.bgLegSynth);
650 pass.dispatchWorkgroups(Math.ceil(this.legLat / WG_SYNTH), mmax + 1);
651 break;
652 case 'fourSynth':
653 pass.setPipeline(this.pipeFourSynth);
654 pass.setBindGroup(0, this.bgFourSynth);
655 if (fft) pass.dispatchWorkgroups(nlat);
656 else pass.dispatchWorkgroups(Math.ceil(nphi / 64), nlat);
657 break;
658 case 'fourAnalys':
659 pass.setPipeline(this.pipeFourAnalys);
660 pass.setBindGroup(0, this.bgFourAnalys);
661 if (fft) pass.dispatchWorkgroups(nlat);
662 else pass.dispatchWorkgroups(Math.ceil((mmax + 1) / 64), nlat);
663 break;
664 case 'legAnalys':
665 pass.setPipeline(this.pipeLegAnalys);
666 pass.setBindGroup(0, this.bgLegAnalys);
667 pass.dispatchWorkgroups(mmax + 1);
668 break;
669 }
670 pass.end();
671 }
673 /** Record the analysis (spatial spatBuf -> spectral qlmOut) into an encoder. */
674 encodeAnalys(encoder: GPUCommandEncoder): void {
675 const pass = encoder.beginComputePass({ label: 'sht-analys' });
676 this.encodeAnalysInto(pass, { bgLeg: this.bgLegAnalys, bgFour: this.bgFourAnalys });
677 pass.end();
678 }
680 /**
681 * Spectral -> spatial. qlm: interleaved [re, im], SHTNS LM ordering,
682 * length 2*nlm. Returns the spatial field, length nlat*nphi.
683 */
684 async synth(qlm: Float32Array): Promise<Float32Array> {
685 const { nlat, nphi } = this.cfg;
686 if (qlm.length !== 2 * this.nlm) throw new Error(`qlm must have length ${2 * this.nlm}`);
687 this.device.queue.writeBuffer(this.qlmIn, 0, qlm as Float32Array<ArrayBuffer>);
688 const enc = this.device.createCommandEncoder();
689 this.encodeSynth(enc);
690 enc.copyBufferToBuffer(this.spatBuf, 0, this.stageSpat, 0, 4 * nlat * nphi);
691 this.device.queue.submit([enc.finish()]);
692 await this.stageSpat.mapAsync(GPUMapMode.READ);
693 const out = new Float32Array(this.stageSpat.getMappedRange().slice(0));
694 this.stageSpat.unmap();
695 return out;
696 }
698 /**
699 * Spectral -> spatial, with the coefficients read from a caller-owned GPU
700 * buffer (interleaved [re, im], 8*nlm bytes, COPY_SRC) instead of uploaded
701 * from the CPU. This is how a field already on the device — a model's
702 * spectral state — is evaluated on this plan's grid, e.g. a finer display
703 * grid than the one the coefficients were produced on.
704 */
705 async synthFrom(qlmSrc: GPUBuffer): Promise<Float32Array> {
706 const { nlat, nphi } = this.cfg;
707 const enc = this.device.createCommandEncoder({ label: 'sht-synth-from' });
708 enc.copyBufferToBuffer(qlmSrc, 0, this.qlmIn, 0, 8 * this.nlm);
709 this.encodeSynth(enc);
710 enc.copyBufferToBuffer(this.spatBuf, 0, this.stageSpat, 0, 4 * nlat * nphi);
711 this.device.queue.submit([enc.finish()]);
712 await this.stageSpat.mapAsync(GPUMapMode.READ);
713 const out = new Float32Array(this.stageSpat.getMappedRange().slice(0));
714 this.stageSpat.unmap();
715 return out;
716 }
718 /** Spatial -> spectral. spat: length nlat*nphi. Returns interleaved qlm, length 2*nlm. */
719 async analys(spat: Float32Array): Promise<Float32Array> {
720 const { nlat, nphi } = this.cfg;
721 if (spat.length !== nlat * nphi) throw new Error(`spat must have length ${nlat * nphi}`);
722 this.device.queue.writeBuffer(this.spatBuf, 0, spat as Float32Array<ArrayBuffer>);
723 const enc = this.device.createCommandEncoder();
724 this.encodeAnalys(enc);
725 enc.copyBufferToBuffer(this.qlmOut, 0, this.stageQ, 0, 8 * this.nlm);
726 this.device.queue.submit([enc.finish()]);
727 await this.stageQ.mapAsync(GPUMapMode.READ);
728 const out = new Float32Array(this.stageQ.getMappedRange().slice(0));
729 this.stageQ.unmap();
730 return out;
731 }
733 destroy(): void {
734 for (const b of [
735 this.bufAb, this.bufAmm, this.bufCtstw, this.bufTrig, this.qlmIn, this.qlmOut,
a4fee9cBatch independent transforms through one Legendre dispatchDan Fortunato 736 this.fmBuf, this.spatBuf, this.stageSpat, this.stageQ, this.fmArena,
738 }
739}
741/** Best-effort human-readable adapter name, so it is clear which GPU (or
742 * software rasterizer) is actually running the transforms. */
743export async function describeAdapter(device: GPUDevice): Promise<string> {
744 const fmt = (info: GPUAdapterInfo | undefined): string => {
745 if (!info) return '';
746 const parts = [info.description, info.device, info.vendor].filter(
747 (s): s is string => !!s && s.length > 0,
748 );
749 const name = parts[0] ?? '';
750 return info.architecture && !name.includes(info.architecture)
751 ? `${name} (${info.architecture})`.trim()
752 : name;
753 };
754 const own = fmt((device as GPUDevice & { adapterInfo?: GPUAdapterInfo }).adapterInfo);
755 if (own) return own;
756 try {
757 const adapter = await navigator.gpu.requestAdapter();
758 return fmt(adapter?.info);
759 } catch {
760 return '';
761 }
762}
764/** Request an adapter/device suitable for the transforms. */
765export async function requestShtDevice(): Promise<GPUDevice> {
766 if (!navigator.gpu) throw new Error('WebGPU is not available in this browser');
767 const adapter = await navigator.gpu.requestAdapter();
768 if (!adapter) throw new Error('No WebGPU adapter available');
769 // ask for a larger workgroup storage if the adapter offers it (bigger FFTs)
770 const wgStorage = Math.min(adapter.limits.maxComputeWorkgroupStorageSize, 32768);
771 // `subgroups` lets the analysis reduction use subgroupAdd instead of a
772 // shared-memory tree (2 barriers per l-pair instead of 1 + log2(wgAnalys)).
773 // Optional: ShtPlan falls back to the tree when it is not available.
774 const features: GPUFeatureName[] = [];
775 if (adapter.features.has('subgroups')) features.push('subgroups');
776 // timestamp-query is only used by the profiling scripts, but it has to be
777 // requested at device creation, and asking costs nothing when unused.
778 if (adapter.features.has('timestamp-query')) features.push('timestamp-query');
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 779 // The seed field's mode table is the one buffer whose size is not fixed by
780 // the grid — it grows with how fine a wavelength is asked for
781 // (src/mgpu/randnfun3.ts), and a browser's default 128 MB storage-buffer
782 // limit is well below what the adapter will actually give. Ask for the
783 // adapter's own maximum so the wavelength is limited by the hardware rather
784 // than by a default.
785 const maxStorage = adapter.limits.maxStorageBufferBindingSize;
786 const maxBuffer = adapter.limits.maxBufferSize;
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 787 return adapter.requestDevice({
788 requiredFeatures: features,
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 789 requiredLimits: {
790 maxComputeWorkgroupStorageSize: wgStorage,
791 maxStorageBufferBindingSize: maxStorage,
792 maxBufferSize: maxBuffer,
793 },
795}