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.
13 */
14import type { ShtPlan, ShtBinding } from './sht.ts';
15import { derivCoeffs } from './derivCoeffs.ts';
16import { dthetaShuffleWGSL, dphiShuffleWGSL, divideSinThetaWGSL } from './wgsl/deriv.ts';
18const WG = 64;
20async function makePipeline(
21 device: GPUDevice,
22 code: string,
23 entryPoint: string,
24): Promise<GPUComputePipeline> {
25 device.pushErrorScope('validation');
26 const module = device.createShaderModule({ code, label: entryPoint });
27 const info = await module.getCompilationInfo();
28 const errors = info.messages.filter((m) => m.type === 'error');
29 if (errors.length) {
30 throw new Error(
31 `WGSL compile error in ${entryPoint}:\n` +
32 errors.map((e) => ` ${e.lineNum}:${e.linePos} ${e.message}`).join('\n'),
33 );
34 }
35 const pipeline = await device.createComputePipelineAsync({
36 layout: 'auto',
37 compute: { module, entryPoint },
38 label: entryPoint,
39 });
40 const err = await device.popErrorScope();
41 if (err) throw new Error(`pipeline ${entryPoint}: ${err.message}`);
42 return pipeline;
43}
45/** Bindings for one dtheta/dphi call against caller-supplied buffers. */
46export interface DerivBinding {
47 readonly shuffle: GPUBindGroup;
48 readonly sht: ShtBinding;
49 /** Only present for dtheta: the post-synthesis divide by sin(theta). */
50 readonly divide?: GPUBindGroup;
51}
53export class DerivPlan {
54 private device: GPUDevice;
55 private sht: ShtPlan;
56 private nlm: number;
57 private npts: number;
59 private bufAPlus!: GPUBuffer;
60 private bufAMinus!: GPUBuffer;
61 private bufMOf!: GPUBuffer;
62 private bufSinTheta!: GPUBuffer;
63 /** Scratch coefficient buffer for the shuffled input to synth -- shared
64 * sequentially like ShtPlan's fmBuf, since ops within one pass execute
65 * in submission order. */
66 private scratch!: GPUBuffer;
68 private pipeDtheta!: GPUComputePipeline;
69 private pipeDphi!: GPUComputePipeline;
70 private pipeDivide!: GPUComputePipeline;
72 private constructor(device: GPUDevice, sht: ShtPlan) {
73 this.device = device;
74 this.sht = sht;
75 this.nlm = sht.nlm;
76 this.npts = sht.cfg.nlat * sht.cfg.nphi;
77 }
79 static async create(device: GPUDevice, sht: ShtPlan): Promise<DerivPlan> {
80 const plan = new DerivPlan(device, sht);
81 await plan.init();
82 return plan;
83 }
85 private async init(): Promise<void> {
86 const { nlat, nphi } = this.sht.cfg;
87 const dev = this.device;
89 const { aPlus, aMinus, mOf } = derivCoeffs(this.sht.cfg.lmax, this.sht.cfg.mmax);
90 const sinTheta = new Float32Array(nlat);
91 for (let i = 0; i < nlat; i++) {
92 const ct = this.sht.cosTheta[i];
93 sinTheta[i] = Math.sqrt(Math.max(0, 1 - ct * ct));
94 }
96 const mk = (label: string, size: number, usage = GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST) =>
97 dev.createBuffer({ label, size, usage });
98 this.bufAPlus = mk('deriv-aplus', 4 * this.nlm);
99 this.bufAMinus = mk('deriv-aminus', 4 * this.nlm);
100 this.bufMOf = mk('deriv-mof', 4 * this.nlm);
101 this.bufSinTheta = mk('deriv-sintheta', 4 * nlat);
102 this.scratch = mk(
103 'deriv-scratch',
104 8 * this.nlm,
105 GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
106 );
108 dev.queue.writeBuffer(this.bufAPlus, 0, new Float32Array(aPlus));
109 dev.queue.writeBuffer(this.bufAMinus, 0, new Float32Array(aMinus));
110 dev.queue.writeBuffer(this.bufMOf, 0, mOf as Uint32Array<ArrayBuffer>);
111 dev.queue.writeBuffer(this.bufSinTheta, 0, sinTheta);
113 const [pDtheta, pDphi, pDivide] = await Promise.all([
114 makePipeline(dev, dthetaShuffleWGSL({ nlm: this.nlm }), 'dtheta_shuffle'),
115 makePipeline(dev, dphiShuffleWGSL({ nlm: this.nlm }), 'dphi_shuffle'),
116 makePipeline(dev, divideSinThetaWGSL({ nlat, nphi }), 'divide_sin_theta'),
117 ]);
118 this.pipeDtheta = pDtheta;
119 this.pipeDphi = pDphi;
120 this.pipeDivide = pDivide;
121 }
123 /** Bindings for dtheta(qlmIn) -> spatOut, against caller-owned buffers. */
124 createDthetaBinding(qlmIn: GPUBuffer, spatOut: GPUBuffer): DerivBinding {
125 const shuffle = this.device.createBindGroup({
126 layout: this.pipeDtheta.getBindGroupLayout(0),
127 entries: [
128 { binding: 0, resource: { buffer: this.bufAPlus } },
129 { binding: 1, resource: { buffer: this.bufAMinus } },
130 { binding: 2, resource: { buffer: qlmIn } },
131 { binding: 3, resource: { buffer: this.scratch } },
132 ],
133 });
134 const sht = this.sht.createSynthBinding(this.scratch, spatOut);
135 const divide = this.device.createBindGroup({
136 layout: this.pipeDivide.getBindGroupLayout(0),
137 entries: [
138 { binding: 0, resource: { buffer: this.bufSinTheta } },
139 { binding: 1, resource: { buffer: spatOut } },
140 ],
141 });
142 return { shuffle, sht, divide };
143 }
145 /** Bindings for dphi(qlmIn) -> spatOut, against caller-owned buffers. */
146 createDphiBinding(qlmIn: GPUBuffer, spatOut: GPUBuffer): DerivBinding {
147 const shuffle = this.device.createBindGroup({
148 layout: this.pipeDphi.getBindGroupLayout(0),
149 entries: [
150 { binding: 0, resource: { buffer: this.bufMOf } },
151 { binding: 1, resource: { buffer: qlmIn } },
152 { binding: 2, resource: { buffer: this.scratch } },
153 ],
154 });
155 const sht = this.sht.createSynthBinding(this.scratch, spatOut);
156 return { shuffle, sht };
157 }
159 /** Record dtheta into an existing compute pass. */
160 encodeDthetaInto(pass: GPUComputePassEncoder, b: DerivBinding): void {
161 pass.setPipeline(this.pipeDtheta);
162 pass.setBindGroup(0, b.shuffle);
163 pass.dispatchWorkgroups(Math.ceil(this.nlm / WG));
164 this.sht.encodeSynthInto(pass, b.sht);
165 pass.setPipeline(this.pipeDivide);
166 pass.setBindGroup(0, b.divide!);
167 pass.dispatchWorkgroups(Math.ceil(this.npts / WG));
168 }
170 /** Record dphi into an existing compute pass. */
171 encodeDphiInto(pass: GPUComputePassEncoder, b: DerivBinding): void {
172 pass.setPipeline(this.pipeDphi);
173 pass.setBindGroup(0, b.shuffle);
174 pass.dispatchWorkgroups(Math.ceil(this.nlm / WG));
175 this.sht.encodeSynthInto(pass, b.sht);
176 }
178 /** CPU convenience: qlm (interleaved [re,im], length 2*nlm) -> grid field. */
179 async dtheta(qlm: Float32Array): Promise<Float32Array> {
180 return this.#runToGrid(qlm, true);
181 }
183 /** CPU convenience: qlm (interleaved [re,im], length 2*nlm) -> grid field. */
184 async dphi(qlm: Float32Array): Promise<Float32Array> {
185 return this.#runToGrid(qlm, false);
186 }
188 async #runToGrid(qlm: Float32Array, withDivide: boolean): Promise<Float32Array> {
189 if (qlm.length !== 2 * this.nlm) throw new Error(`qlm must have length ${2 * this.nlm}`);
190 const dev = this.device;
191 const qlmIn = dev.createBuffer({
192 label: 'deriv-qlm-in',
193 size: 8 * this.nlm,
194 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
195 });
196 const spatOut = dev.createBuffer({
197 label: 'deriv-spat-out',
198 size: 4 * this.npts,
199 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
200 });
201 const stage = dev.createBuffer({
202 label: 'deriv-stage',
203 size: 4 * this.npts,
204 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
205 });
206 try {
207 dev.queue.writeBuffer(qlmIn, 0, qlm as Float32Array<ArrayBuffer>);
208 const binding = withDivide
209 ? this.createDthetaBinding(qlmIn, spatOut)
210 : this.createDphiBinding(qlmIn, spatOut);
211 const enc = dev.createCommandEncoder({ label: 'deriv-run' });
212 const pass = enc.beginComputePass({ label: 'deriv-run' });
213 if (withDivide) this.encodeDthetaInto(pass, binding);
214 else this.encodeDphiInto(pass, binding);
215 pass.end();
216 enc.copyBufferToBuffer(spatOut, 0, stage, 0, 4 * this.npts);
217 dev.queue.submit([enc.finish()]);
218 await stage.mapAsync(GPUMapMode.READ);
219 const out = new Float32Array(stage.getMappedRange().slice(0));
220 stage.unmap();
221 return out;
222 } finally {
223 qlmIn.destroy();
224 spatOut.destroy();
225 stage.destroy();
226 }
227 }
229 destroy(): void {
230 for (const b of [
231 this.bufAPlus, this.bufAMinus, this.bufMOf, this.bufSinTheta, this.scratch,
232 ]) b?.destroy();
233 }
234}