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