1/**
2 * The WGSL spherical-harmonic transforms against the f64 CPU reference.
3 *
4 * This is the one place a second implementation is still the right oracle: the
5 * transforms are vendored shtns-webgpu, and `src/sht/reference.ts` is its direct-
6 * summation f64 twin. Everything above them (the .m models) is checked against
7 * closed-form answers instead — see analyticChecks.ts.
8 */
9import { ShtPlan } from '../src/sht/sht.ts';
10import { ShtReference, randomSpectrum } from '../src/sht/reference.ts';
11import { DerivPlan } from '../src/sht/deriv.ts';
12import { gridForLmax } from '../src/sht/layout.ts';
13import type { Check, Log } from './analyticChecks.ts';
15function relL2(a: ArrayLike<number>, b: ArrayLike<number>): number {
16 let num = 0;
17 let den = 0;
18 for (let i = 0; i < a.length; i++) {
19 const d = a[i] - b[i];
20 num += d * d;
21 den += b[i] * b[i];
22 }
23 return Math.sqrt(num / Math.max(den, 1e-300));
24}
26export async function transformChecks(
27 device: GPUDevice,
28 check: Check,
29 _log: Log,
30): Promise<void> {
31 const lmax = 31;
32 const { nlat, nphi } = gridForLmax(lmax, 1);
33 const cfg = { lmax, mmax: lmax, nlat, nphi };
35 const plan = await ShtPlan.create(device, cfg);
36 const ref = new ShtReference(cfg);
38 const q = randomSpectrum(cfg, 42);
39 const q64 = new Float64Array(q);
41 const spatGpu = await plan.synth(new Float32Array(q64));
42 const spatCpu = ref.synth(q64);
43 const errSynth = relL2(spatGpu, spatCpu);
45 const qGpu = await plan.analys(new Float32Array(spatCpu));
46 const qCpu = ref.analys(new Float64Array(spatCpu));
47 const errAnalys = relL2(qGpu, qCpu);
49 check(
50 'transforms: WGSL fp32 vs f64 CPU reference',
51 errSynth < 1e-4 && errAnalys < 1e-4,
52 `synth ${errSynth.toExponential(2)}, analys ${errAnalys.toExponential(2)}`,
53 );
55 // ---- f64 reference dtheta/dphi vs an independent closed form -----------
56 // x(theta,phi) = sin(theta)*cos(phi) is exactly degree 1, so quadrature
57 // recovers it to f64 round-off; comparing its dtheta/dphi against the
58 // grid-space analytic derivatives (not derived from the same recurrence
59 // being tested) catches a sign or indexing error the random-spectrum check
60 // below, which compares two implementations of the same formula, would not.
61 {
62 const x = new Float64Array(nlat * nphi);
63 for (let i = 0; i < nlat; i++) {
64 const st = ref.st[i];
65 for (let j = 0; j < nphi; j++) {
66 const phi = (2 * Math.PI * j) / nphi;
67 x[i * nphi + j] = st * Math.cos(phi);
68 }
69 }
70 const X = ref.analys(x);
71 const dThetaX = ref.dtheta(X);
72 const dPhiX = ref.dphi(X);
74 let errNum = 0;
75 let norm = 0;
76 for (let i = 0; i < nlat; i++) {
77 const ct = ref.ct[i];
78 const st = ref.st[i];
79 for (let j = 0; j < nphi; j++) {
80 const phi = (2 * Math.PI * j) / nphi;
81 const k = i * nphi + j;
82 const wantTheta = ct * Math.cos(phi);
83 const wantPhi = -st * Math.sin(phi);
84 errNum += (dThetaX[k] - wantTheta) ** 2 + (dPhiX[k] - wantPhi) ** 2;
85 norm += wantTheta * wantTheta + wantPhi * wantPhi;
86 }
87 }
88 const relErr = Math.sqrt(errNum / Math.max(norm, 1e-300));
89 check(
90 'deriv: f64 reference dtheta/dphi match the closed form on x = sin(theta)cos(phi)',
91 relErr < 1e-6,
92 `rel L2 error ${relErr.toExponential(2)}`,
93 );
94 }
96 // ---- WGSL fp32 dtheta/dphi vs the (now closed-form-verified) f64 reference
97 {
98 const deriv = await DerivPlan.create(device, plan);
100 const dThetaGpu = await deriv.dtheta(new Float32Array(q64));
101 const dThetaCpu = ref.dtheta(q64);
102 const errDtheta = relL2(dThetaGpu, dThetaCpu);
104 const dPhiGpu = await deriv.dphi(new Float32Array(q64));
105 const dPhiCpu = ref.dphi(q64);
106 const errDphi = relL2(dPhiGpu, dPhiCpu);
108 check(
109 'deriv: WGSL fp32 dtheta/dphi vs f64 CPU reference',
110 errDtheta < 1e-4 && errDphi < 1e-4,
111 `dtheta ${errDtheta.toExponential(2)}, dphi ${errDphi.toExponential(2)}`,
112 );
114 // The undivided theta derivative sin(theta)*dtheta(u) — the flux-form
115 // Laplace-Beltrami scheme's step 1 and the flux-metric precompute's
116 // input — is the same shuffle+synthesis with the divide skipped, so it
117 // gets the same oracle.
118 const sinDthetaGpu = await deriv.sinDtheta(new Float32Array(q64));
119 const sinDthetaCpu = ref.sinDtheta(q64);
120 const errSinDtheta = relL2(sinDthetaGpu, sinDthetaCpu);
121 check(
122 'deriv: WGSL fp32 sinDtheta (undivided) vs f64 CPU reference',
123 errSinDtheta < 1e-4,
124 `sinDtheta ${errSinDtheta.toExponential(2)}`,
125 );
126 deriv.destroy();
127 }
129 // ---- batched transforms reproduce the scalar transforms ------------------
130 // A batch walks the Legendre recurrence once for K fields with per-lane
131 // arithmetic textually identical to the scalar kernel's, so each lane must
132 // agree with the scalar path to shader-compiler latitude (FMA contraction
133 // may differ between the two modules; nothing else may).
134 {
135 const { nlat: gl, nphi: gp } = plan.cfg;
136 const npts = gl * gp;
137 const sizes = [];
138 for (let k = 2; k <= plan.batchK; k += 2) sizes.push(k);
139 check(
140 'batch: plan compiled batched pipelines',
141 plan.batchK >= 2,
142 `batchK = ${plan.batchK} (${sizes.map((s) => `x${s}`).join(', ') || 'none'})`,
143 );
144 for (const K of sizes) {
145 const qs = Array.from({ length: K }, (_, k) => randomSpectrum(cfg, 1000 + k));
146 const qBufs = qs.map((q, k) => {
147 const b = device.createBuffer({
148 label: `batch-test-q${k}`,
149 size: 8 * plan.nlm,
150 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
151 });
152 device.queue.writeBuffer(b, 0, q as Float32Array<ArrayBuffer>);
153 return b;
154 });
155 const spatBufs = qs.map((_, k) =>
156 device.createBuffer({
157 label: `batch-test-spat${k}`,
158 size: 4 * npts,
159 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
160 }),
161 );
162 const qOutBufs = qs.map((_, k) =>
163 device.createBuffer({
164 label: `batch-test-qout${k}`,
165 size: 8 * plan.nlm,
166 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
167 }),
168 );
169 const stage = device.createBuffer({
170 label: 'batch-test-stage',
171 size: K * (4 * npts + 8 * plan.nlm),
172 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
173 });
175 // One pass: batched synthesis of all K, then batched analysis back.
176 const synthB = plan.createSynthBatchBinding(
177 qs.map((_, k) => ({ qlmIn: qBufs[k], spatOut: spatBufs[k] })),
178 );
179 const analysB = plan.createAnalysBatchBinding(
180 qs.map((_, k) => ({ spatIn: spatBufs[k], qlmOut: qOutBufs[k] })),
181 );
182 const enc = device.createCommandEncoder({ label: 'batch-test' });
183 const pass = enc.beginComputePass();
184 plan.encodeSynthBatchInto(pass, synthB);
185 plan.encodeAnalysBatchInto(pass, analysB);
186 pass.end();
187 for (let k = 0; k < K; k++) {
188 enc.copyBufferToBuffer(spatBufs[k], 0, stage, k * 4 * npts, 4 * npts);
189 enc.copyBufferToBuffer(qOutBufs[k], 0, stage, K * 4 * npts + k * 8 * plan.nlm, 8 * plan.nlm);
190 }
191 device.queue.submit([enc.finish()]);
192 await stage.mapAsync(GPUMapMode.READ);
193 const raw = new Float32Array(stage.getMappedRange().slice(0));
194 stage.unmap();
196 let worstSynth = 0;
197 let worstAnalys = 0;
198 for (let k = 0; k < K; k++) {
199 const spatLane = raw.subarray(k * npts, (k + 1) * npts);
200 const qLane = raw.subarray(K * npts + k * 2 * plan.nlm, K * npts + (k + 1) * 2 * plan.nlm);
201 const spatScalar = await plan.synth(qs[k]);
202 const qScalar = await plan.analys(spatScalar);
203 worstSynth = Math.max(worstSynth, relL2(spatLane, spatScalar));
204 worstAnalys = Math.max(worstAnalys, relL2(qLane, qScalar));
205 }
206 check(
207 `batch: x${K} lanes match the scalar transforms`,
208 worstSynth < 1e-6 && worstAnalys < 1e-6,
209 `synth ${worstSynth.toExponential(2)}, analys ${worstAnalys.toExponential(2)} ` +
210 `across ${K} lanes`,
211 );
212 for (const b of [...qBufs, ...spatBufs, ...qOutBufs, stage]) b.destroy();
213 }
214 }
216 plan.destroy();
217}