8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 1/**
2 * First derivatives of a scalar field, coefficients -> grid: dtheta and dphi
3 * (evolving_surface/notes/algos.tex Algorithm 1, theta/phi branches only --
4 * the Laplace-Beltrami operator built on these never needs the second-
5 * derivative/curvature branches, so they are not ported).
6 *
7 * Both derivatives start with a shuffle in coefficient space (the theta
8 * branch's +-1 index gather via the alpha recurrence, the phi branch's i*m
9 * row-swap) and then reuse the *existing* Legendre+Fourier synthesis
10 * pipeline (ShtPlan.createSynthBinding/encodeSynthInto) unchanged -- neither
11 * derivative touches the Legendre recurrence stage itself. dtheta
12 * additionally divides by sin(theta) on the grid afterwards.
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 13 *
14 * The two shuffles are also exposed on their own, coefficients -> coefficients
15 * (`dthetac`, `dphic`), because the six-transform Laplace-Beltrami operator of
16 * docs/reduced-transforms.md needs them apart from a
17 * synthesis, and needs them twice: once on the field (steps 1-2) and once on
18 * the two fluxes (step 5, which is the *same* alpha^+/alpha^- gather, not its
19 * transpose). Everything above is then a composition of them:
20 *
21 * dphi(U) == synth(dphic(U))
22 * dtheta(U) == synth(dthetac(U)) / sin(theta)
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 23 */
24import type { ShtPlan, ShtBinding } from './sht.ts';
25import { derivCoeffs } from './derivCoeffs.ts';
26import { dthetaShuffleWGSL, dphiShuffleWGSL, divideSinThetaWGSL } from './wgsl/deriv.ts';
28const WG = 64;
30async function makePipeline(
31 device: GPUDevice,
32 code: string,
33 entryPoint: string,
34): Promise<GPUComputePipeline> {
35 device.pushErrorScope('validation');
36 const module = device.createShaderModule({ code, label: entryPoint });
37 const info = await module.getCompilationInfo();
38 const errors = info.messages.filter((m) => m.type === 'error');
39 if (errors.length) {
40 throw new Error(
41 `WGSL compile error in ${entryPoint}:\n` +
42 errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n'),
43 );
44 }
45 const pipeline = await device.createComputePipelineAsync({
46 layout: 'auto',
47 compute: { module, entryPoint },
48 label: entryPoint,
49 });
50 const err = await device.popErrorScope();
51 if (err) throw new Error(`pipeline ${entryPoint}: ${err.message}`);
52 return pipeline;
53}
55/** Bindings for one dtheta/dphi call against caller-supplied buffers. */
56export interface DerivBinding {
57 readonly shuffle: GPUBindGroup;
58 readonly sht: ShtBinding;
59 /** Only present for dtheta: the post-synthesis divide by sin(theta). */
60 readonly divide?: GPUBindGroup;
61}
63export class DerivPlan {
64 private device: GPUDevice;
65 private sht: ShtPlan;
66 private nlm: number;
67 private npts: number;
69 private bufAPlus!: GPUBuffer;
70 private bufAMinus!: GPUBuffer;
71 private bufMOf!: GPUBuffer;
72 private bufSinTheta!: GPUBuffer;
73 /** Scratch coefficient buffer for the shuffled input to synth -- shared
74 * sequentially like ShtPlan's fmBuf, since ops within one pass execute
75 * in submission order. */
76 private scratch!: GPUBuffer;
78 private pipeDtheta!: GPUComputePipeline;
79 private pipeDphi!: GPUComputePipeline;
80 private pipeDivide!: GPUComputePipeline;
82 private constructor(device: GPUDevice, sht: ShtPlan) {
83 this.device = device;
84 this.sht = sht;
85 this.nlm = sht.nlm;
86 this.npts = sht.cfg.nlat * sht.cfg.nphi;
87 }
89 static async create(device: GPUDevice, sht: ShtPlan): Promise<DerivPlan> {
90 const plan = new DerivPlan(device, sht);
91 await plan.init();
92 return plan;
93 }
95 private async init(): Promise<void> {
96 const { nlat, nphi } = this.sht.cfg;
97 const dev = this.device;
99 const { aPlus, aMinus, mOf } = derivCoeffs(this.sht.cfg.lmax, this.sht.cfg.mmax);
100 const sinTheta = new Float32Array(nlat);
101 for (let i = 0; i < nlat; i++) {
102 const ct = this.sht.cosTheta[i];
103 sinTheta[i] = Math.sqrt(Math.max(0, 1 - ct * ct));
104 }
106 const mk = (label: string, size: number, usage = GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST) =>
107 dev.createBuffer({ label, size, usage });
108 this.bufAPlus = mk('deriv-aplus', 4 * this.nlm);
109 this.bufAMinus = mk('deriv-aminus', 4 * this.nlm);
110 this.bufMOf = mk('deriv-mof', 4 * this.nlm);
111 this.bufSinTheta = mk('deriv-sintheta', 4 * nlat);
112 this.scratch = mk(
113 'deriv-scratch',
114 8 * this.nlm,
115 GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
116 );
118 dev.queue.writeBuffer(this.bufAPlus, 0, new Float32Array(aPlus));
119 dev.queue.writeBuffer(this.bufAMinus, 0, new Float32Array(aMinus));
120 dev.queue.writeBuffer(this.bufMOf, 0, mOf as Uint32Array<ArrayBuffer>);
121 dev.queue.writeBuffer(this.bufSinTheta, 0, sinTheta);
123 const [pDtheta, pDphi, pDivide] = await Promise.all([
124 makePipeline(dev, dthetaShuffleWGSL({ nlm: this.nlm }), 'dtheta_shuffle'),
125 makePipeline(dev, dphiShuffleWGSL({ nlm: this.nlm }), 'dphi_shuffle'),
126 makePipeline(dev, divideSinThetaWGSL({ nlat, nphi }), 'divide_sin_theta'),
127 ]);
128 this.pipeDtheta = pDtheta;
129 this.pipeDphi = pDphi;
130 this.pipeDivide = pDivide;
131 }
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 133 /**
134 * Bind group for the coefficient-space half of dtheta on its own:
135 * v_l^m = alpha^+(l-1,m) u_{l-1}^m + alpha^-(l+1,m) u_{l+1}^m, the
136 * coefficients of sin(theta) * dtheta(u). Input and output must be
137 * different buffers -- WebGPU forbids binding one buffer as both readable
138 * and writable storage in a dispatch, and the gather reads l+-1 anyway.
139 */
140 createDthetacBinding(qlmIn: GPUBuffer, qlmOut: GPUBuffer): GPUBindGroup {
141 return this.device.createBindGroup({
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 142 layout: this.pipeDtheta.getBindGroupLayout(0),
143 entries: [
144 { binding: 0, resource: { buffer: this.bufAPlus } },
145 { binding: 1, resource: { buffer: this.bufAMinus } },
146 { binding: 2, resource: { buffer: qlmIn } },
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 147 { binding: 3, resource: { buffer: qlmOut } },
148 ],
149 });
150 }
152 /** Bind group for the coefficient-space half of dphi on its own:
153 * (dphi u)_l^m = i*m*u_l^m. Same buffer restriction as dthetac. */
154 createDphicBinding(qlmIn: GPUBuffer, qlmOut: GPUBuffer): GPUBindGroup {
155 return this.device.createBindGroup({
156 layout: this.pipeDphi.getBindGroupLayout(0),
157 entries: [
158 { binding: 0, resource: { buffer: this.bufMOf } },
159 { binding: 1, resource: { buffer: qlmIn } },
160 { binding: 2, resource: { buffer: qlmOut } },
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 161 ],
162 });
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 163 }
165 /** Record the bare alpha^+/alpha^- shift into an existing compute pass. */
166 encodeDthetacInto(pass: GPUComputePassEncoder, bindGroup: GPUBindGroup): void {
167 pass.setPipeline(this.pipeDtheta);
168 pass.setBindGroup(0, bindGroup);
169 pass.dispatchWorkgroups(Math.ceil(this.nlm / WG));
170 }
172 /** Record the bare i*m multiply into an existing compute pass. */
173 encodeDphicInto(pass: GPUComputePassEncoder, bindGroup: GPUBindGroup): void {
174 pass.setPipeline(this.pipeDphi);
175 pass.setBindGroup(0, bindGroup);
176 pass.dispatchWorkgroups(Math.ceil(this.nlm / WG));
177 }
179 /** Bindings for dtheta(qlmIn) -> spatOut, against caller-owned buffers. */
180 createDthetaBinding(qlmIn: GPUBuffer, spatOut: GPUBuffer): DerivBinding {
181 const shuffle = this.createDthetacBinding(qlmIn, this.scratch);
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 182 const sht = this.sht.createSynthBinding(this.scratch, spatOut);
183 const divide = this.device.createBindGroup({
184 layout: this.pipeDivide.getBindGroupLayout(0),
185 entries: [
186 { binding: 0, resource: { buffer: this.bufSinTheta } },
187 { binding: 1, resource: { buffer: spatOut } },
188 ],
189 });
190 return { shuffle, sht, divide };
191 }
193 /** Bindings for dphi(qlmIn) -> spatOut, against caller-owned buffers. */
194 createDphiBinding(qlmIn: GPUBuffer, spatOut: GPUBuffer): DerivBinding {
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 195 const shuffle = this.createDphicBinding(qlmIn, this.scratch);
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 196 const sht = this.sht.createSynthBinding(this.scratch, spatOut);
197 return { shuffle, sht };
198 }
200 /** Record dtheta into an existing compute pass. */
201 encodeDthetaInto(pass: GPUComputePassEncoder, b: DerivBinding): void {
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 202 this.encodeSinDthetaInto(pass, b);
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 203 pass.setPipeline(this.pipeDivide);
204 pass.setBindGroup(0, b.divide!);
205 pass.dispatchWorkgroups(Math.ceil(this.npts / WG));
206 }
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 208 /** Record dtheta *without* its final division: the grid values of
209 * sin(theta) * dtheta(u), which unlike dtheta(u) itself is a smooth
210 * function on the sphere. Takes a dtheta binding and simply stops early. */
211 encodeSinDthetaInto(pass: GPUComputePassEncoder, b: DerivBinding): void {
212 this.encodeDthetacInto(pass, b.shuffle);
213 this.sht.encodeSynthInto(pass, b.sht);
214 }
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 216 /** Record dphi into an existing compute pass. */
217 encodeDphiInto(pass: GPUComputePassEncoder, b: DerivBinding): void {
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 218 this.encodeDphicInto(pass, b.shuffle);
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 219 this.sht.encodeSynthInto(pass, b.sht);
220 }
222 /** CPU convenience: qlm (interleaved [re,im], length 2*nlm) -> grid field. */
223 async dtheta(qlm: Float32Array): Promise<Float32Array> {
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 224 return this.#runToGrid(qlm, 'dtheta');
225 }
227 /** CPU convenience: sin(theta) * dtheta(u) on the grid, the undivided
228 * synthesis of the alpha shift. What the flux-form metric precompute
229 * (src/geom/metric.ts) is built from. */
230 async sinDtheta(qlm: Float32Array): Promise<Float32Array> {
231 return this.#runToGrid(qlm, 'sinDtheta');
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 232 }
234 /** CPU convenience: qlm (interleaved [re,im], length 2*nlm) -> grid field. */
235 async dphi(qlm: Float32Array): Promise<Float32Array> {
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 236 return this.#runToGrid(qlm, 'dphi');
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 237 }
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 239 async #runToGrid(
240 qlm: Float32Array,
241 mode: 'dtheta' | 'sinDtheta' | 'dphi',
242 ): Promise<Float32Array> {
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 243 if (qlm.length !== 2 * this.nlm) throw new Error(`qlm must have length ${2 * this.nlm}`);
244 const dev = this.device;
245 const qlmIn = dev.createBuffer({
246 label: 'deriv-qlm-in',
247 size: 8 * this.nlm,
248 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
249 });
250 const spatOut = dev.createBuffer({
251 label: 'deriv-spat-out',
252 size: 4 * this.npts,
253 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
254 });
255 const stage = dev.createBuffer({
256 label: 'deriv-stage',
257 size: 4 * this.npts,
258 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
259 });
260 try {
261 dev.queue.writeBuffer(qlmIn, 0, qlm as Float32Array<ArrayBuffer>);
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 262 const binding =
263 mode === 'dphi'
264 ? this.createDphiBinding(qlmIn, spatOut)
265 : this.createDthetaBinding(qlmIn, spatOut);
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 266 const enc = dev.createCommandEncoder({ label: 'deriv-run' });
267 const pass = enc.beginComputePass({ label: 'deriv-run' });
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 268 if (mode === 'dtheta') this.encodeDthetaInto(pass, binding);
269 else if (mode === 'sinDtheta') this.encodeSinDthetaInto(pass, binding);
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 270 else this.encodeDphiInto(pass, binding);
271 pass.end();
272 enc.copyBufferToBuffer(spatOut, 0, stage, 0, 4 * this.npts);
273 dev.queue.submit([enc.finish()]);
274 await stage.mapAsync(GPUMapMode.READ);
275 const out = new Float32Array(stage.getMappedRange().slice(0));
276 stage.unmap();
277 return out;
278 } finally {
279 qlmIn.destroy();
280 spatOut.destroy();
281 stage.destroy();
282 }
283 }
285 destroy(): void {
286 for (const b of [
287 this.bufAPlus, this.bufAMinus, this.bufMOf, this.bufSinTheta, this.scratch,
288 ]) b?.destroy();
289 }
290}