2 * A .m model, compiled and running on the GPU.
3 *
4 * A model file is ordinary MATLAB: it defines an `init` function that builds the
5 * initial spectral state and a `step` function that advances it one timestep.
6 * Each is specialized for the current grid and compiled into a ModelPlan, and
7 * both operate on the same state buffers (see HostBuffers).
8 *
9 * Both functions return the new state followed by the grid fields the app
10 * renders, so their signatures say exactly what they produce:
11 *
12 * function [U, V, u, v] = init(noise, a, b)
13 * function [U, V, u, v] = step(U, V, lam, a, b, D1, D2, dt)
14 *
15 * The host supplies the things that are precomputation rather than algorithm:
16 * the grid, the Laplace-Beltrami eigenvalues, the seeded initial noise, and the
17 * parameter values. Each argument is matched to the .m's declared parameter
18 * name, so the file documents its own interface.
19 */
20import { ShtPlan } from '../sht/sht.ts';
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 21import type { DerivPlan } from '../sht/deriv.ts';
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 22import { lmIndex, type ShtConfig } from '../sht/layout.ts';
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 23import { HostBuffers, ModelPlan, type Randnfun3Lambda } from './plan.ts';
24import { MODE_BUFFER } from './randnfun3.ts';
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 25import { inFunction, inFunctionAsync, inModel } from './errors.ts';
26import { CompiledModel, type Binding } from './compile.ts';
28export interface ModelParams {
29 [key: string]: number;
30}
32export interface GpuModelOptions {
33 device: GPUDevice;
34 sht: ShtPlan;
35 cfg: ShtConfig;
36 /** Model source (.m text). */
37 source: string;
38 /** Parameter names the .m may take as arguments. */
39 paramNames: string[];
40 /** Spectral state names, in order (e.g. ['U', 'V']). */
41 state: string[];
42 /** Grid fields to render, in order (e.g. ['u', 'v']). */
43 view: string[];
44 /**
45 * The surface, as the .m may ask for it: `gx`, `gy`, `gz` are the embedding's
46 * Cartesian coordinates on the grid, and `Gx`, `Gy`, `Gz` the spherical-
47 * harmonic coefficients they were synthesized from. Omitted for a bare unit
48 * sphere, where the .m has no geometry to take.
49 */
50 geometry?: GeometryBuffers;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 51 /** Computes `dtheta`/`dphi` for the .m's surface Laplace-Beltrami
52 * correction. Omitted for a bare unit sphere, same as `geometry`. */
53 deriv?: DerivPlan;
55 * Iterations of the implicit solve the .m's `for` loop runs. A fixed scalar
56 * rather than a tunable one: the loop is unrolled into the op sequence, so
57 * the count is part of what compiles and changing it recompiles.
58 */
59 niter?: number;
60}
62/** Host-supplied surface fields, in the layout the .m sees them. */
63export interface GeometryBuffers {
64 /** Grid coordinates, npts each. */
65 x: Float32Array;
66 y: Float32Array;
67 z: Float32Array;
68 /** Their spherical-harmonic coefficients, 2 x nlm each. */
69 X: Float32Array;
70 Y: Float32Array;
71 Z: Float32Array;
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 72 /** Inverse metric quantities (src/geom/metric.ts), grid space, npts each —
73 * the Algorithm-4 (12-transform) Laplace-Beltrami path. */
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 74 Vtx: Float32Array;
75 Vty: Float32Array;
76 Vtz: Float32Array;
77 Vpx: Float32Array;
78 Vpy: Float32Array;
79 Vpz: Float32Array;
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 80 /** Flux-form metric weights (src/geom/metric.ts computeFluxMetric), grid
81 * space, npts each — the six-transform Laplace-Beltrami scheme of
82 * docs/reduced-transforms.md. */
83 p1: Float32Array;
84 p2: Float32Array;
85 q2: Float32Array;
86 r: Float32Array;
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 87 /** Mean-J preconditioner scale (Geometry.Jhat): folded into every
88 * setParams upload as the 'jhat' uniform, so a .m that takes jhat is
89 * never left with the zero a missing parameter would default to. An
90 * explicit jhat in the params wins (jhat: 1 pins the plain round-sphere
91 * preconditioner, for A/B). */
92 Jhat: number;
95/** Names the .m may take for the grid coordinates and for their coefficients. */
96export const GEOMETRY_GRID_NAMES = ['gx', 'gy', 'gz'] as const;
97export const GEOMETRY_SPECTRAL_NAMES = ['Gx', 'Gy', 'Gz'] as const;
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 98/** Names the .m may take for the inverse metric quantities (Algorithm 4). */
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 99export const METRIC_GRID_NAMES = ['Vtx', 'Vty', 'Vtz', 'Vpx', 'Vpy', 'Vpz'] as const;
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 100/** Names the .m may take for the flux-form metric weights (six-transform
101 * scheme). A model asks for whichever set its loop uses; both are uploaded. */
102export const FLUX_METRIC_GRID_NAMES = ['p1', 'p2', 'q2', 'r'] as const;
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 103
104/** Laplace-Beltrami eigenvalues l(l+1), duplicated across re/im so the array
105 * matches the 2 x nlm spectral layout element for element. */
106export function eigenvalues(cfg: ShtConfig, nlm: number): Float32Array {
107 const lam = new Float32Array(2 * nlm);
108 for (let m = 0; m <= cfg.mmax; m++) {
109 for (let l = m; l <= cfg.lmax; l++) {
110 const i = lmIndex(cfg.lmax, l, m);
111 lam[2 * i] = l * (l + 1);
112 lam[2 * i + 1] = l * (l + 1);
113 }
114 }
115 return lam;
116}
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 118/**
119 * 1 where l < lmax-2, else 0, duplicated across re/im like `lam`. The
120 * theta/phi derivative recurrences (src/sht/derivCoeffs.ts) cannot exactly
121 * represent a derivative at the top two degrees of the band limit, so the
122 * surface Laplace-Beltrami correction filters them out wherever it
123 * re-differentiates a field (evolving_surface/notes/algos.tex Sec 6,
124 * "Miscellaneous implementation details").
125 */
126export function filterMask(cfg: ShtConfig, nlm: number): Float32Array {
127 const filt = new Float32Array(2 * nlm);
128 for (let m = 0; m <= cfg.mmax; m++) {
129 for (let l = m; l <= cfg.lmax; l++) {
130 const i = lmIndex(cfg.lmax, l, m);
131 const keep = l < cfg.lmax - 2 ? 1 : 0;
132 filt[2 * i] = keep;
133 filt[2 * i + 1] = keep;
134 }
135 }
136 return filt;
137}
140 readonly paramNames: string[];
141 readonly state: string[];
142 readonly view: string[];
143 readonly npts: number;
144 readonly nlm: number;
146 #device: GPUDevice;
147 #host: HostBuffers;
148 #initPlan: ModelPlan;
149 #stepPlan: ModelPlan;
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 150 /** Current geometry's mean-J scale; 1 with no geometry (the sphere). */
151 #jhat = 1;
153 /** Scratch holding a copy of the whole spectral state; see snapshotState. */
154 #stash: GPUBuffer;
155 /** Which function wrote the state most recently; see `read`. */
156 #lastRan: 'init' | 'step' = 'init';
157 #stashedRan: 'init' | 'step' = 'init';
158 #destroyed = false;
160 private constructor(init: {
161 device: GPUDevice;
162 host: HostBuffers;
163 initPlan: ModelPlan;
164 stepPlan: ModelPlan;
165 readback: GPUBuffer;
166 stash: GPUBuffer;
167 paramNames: string[];
168 state: string[];
169 view: string[];
170 npts: number;
171 nlm: number;
172 }) {
173 this.#device = init.device;
174 this.#host = init.host;
175 this.#initPlan = init.initPlan;
176 this.#stepPlan = init.stepPlan;
177 this.#readback = init.readback;
178 this.#stash = init.stash;
179 this.paramNames = init.paramNames;
180 this.state = init.state;
181 this.view = init.view;
182 this.npts = init.npts;
183 this.nlm = init.nlm;
184 }
186 static async create(opts: GpuModelOptions): Promise<GpuModel> {
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 187 const { device, sht, cfg, source, paramNames, state, view, geometry, deriv } = opts;
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 188 const npts = cfg.nlat * cfg.nphi;
189 const nlm = sht.nlm;
190 const niter = opts.niter ?? 0;
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 192 // What the .m may ask for by parameter name. Spectral state, the
193 // eigenvalues and the top-mode filter are 2 x nlm; the seeded
194 // perturbation is a grid field.
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 195 const bindings: Record<string, Binding> = {
196 lam: { kind: 'tensor', shape: [2, nlm] },
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 197 filt: { kind: 'tensor', shape: [2, nlm] },
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 198 noise: { kind: 'tensor', shape: [npts, 1] },
199 npts: { kind: 'const', value: npts },
200 nlm: { kind: 'const', value: nlm },
201 niter: { kind: 'const', value: niter },
202 };
203 if (geometry) {
204 for (const g of GEOMETRY_GRID_NAMES) bindings[g] = { kind: 'tensor', shape: [npts, 1] };
205 for (const g of GEOMETRY_SPECTRAL_NAMES) bindings[g] = { kind: 'tensor', shape: [2, nlm] };
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 206 for (const g of METRIC_GRID_NAMES) bindings[g] = { kind: 'tensor', shape: [npts, 1] };
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 207 for (const g of FLUX_METRIC_GRID_NAMES) bindings[g] = { kind: 'tensor', shape: [npts, 1] };
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 208 // Mean-J preconditioner scale (Geometry.Jhat): a uniform, not a const,
209 // so swapping the surface updates it with no recompile. The session
210 // folds the current geometry's value into every setParams call.
211 bindings['jhat'] = { kind: 'param' };
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 212 // The wavelength of the seeded random field (src/mgpu/randnfun3.ts).
213 // A uniform like jhat, not a const: changing it redraws the field
214 // without recompiling the step.
215 bindings['lam3'] = { kind: 'param' };
217 for (const s of state) bindings[s] = { kind: 'tensor', shape: [2, nlm] };
218 for (const p of paramNames) bindings[p] = { kind: 'param' };
220 // Parsing belongs to the file, not to either function.
221 const compiled = inModel(() => new CompiledModel(source, bindings, { npts, nlm }));
222 // Both functions return the new state first, then the rendered grid fields.
223 const nargout = state.length + view.length;
224 const initFn = inFunction('init', () => compiled.specialize('init', nargout));
225 const stepFn = inFunction('step', () => compiled.specialize('step', nargout));
226 compiled.finish();
228 // Only the state outputs feed back into the argument buffers; the grid
229 // fields are read for display and then overwritten next call.
230 const feedback = [...state, ...view.map(() => null)];
232 const host = new HostBuffers(device);
233 // The host owns the state and the inputs it uploads, whether or not a given
234 // function happens to take them as arguments — `init` does not read `U`, but
235 // it writes it, and `step` reads it back.
236 for (const s of state) host.ensure(s, 2 * nlm);
237 host.ensure('lam', 2 * nlm);
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 238 host.ensure('filt', 2 * nlm);
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 239 host.ensure('noise', npts);
240 if (geometry) {
241 for (const g of GEOMETRY_GRID_NAMES) host.ensure(g, npts);
242 for (const g of GEOMETRY_SPECTRAL_NAMES) host.ensure(g, 2 * nlm);
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 243 for (const g of METRIC_GRID_NAMES) host.ensure(g, npts);
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 244 for (const g of FLUX_METRIC_GRID_NAMES) host.ensure(g, npts);
247 const initPlan = await inFunctionAsync('init', () =>
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 248 ModelPlan.create(device, sht, { fn: initFn, feedback }, host, deriv),
250 const stepPlan = await inFunctionAsync('step', () =>
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 251 ModelPlan.create(device, sht, { fn: stepFn, feedback }, host, deriv),
254 host.upload('lam', eigenvalues(cfg, nlm));
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 255 host.upload('filt', filterMask(cfg, nlm));
257 host.upload('gx', geometry.x);
258 host.upload('gy', geometry.y);
259 host.upload('gz', geometry.z);
260 host.upload('Gx', geometry.X);
261 host.upload('Gy', geometry.Y);
262 host.upload('Gz', geometry.Z);
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 263 host.upload('Vtx', geometry.Vtx);
264 host.upload('Vty', geometry.Vty);
265 host.upload('Vtz', geometry.Vtz);
266 host.upload('Vpx', geometry.Vpx);
267 host.upload('Vpy', geometry.Vpy);
268 host.upload('Vpz', geometry.Vpz);
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 269 host.upload('p1', geometry.p1);
270 host.upload('p2', geometry.p2);
271 host.upload('q2', geometry.q2);
272 host.upload('r', geometry.r);
275 const readback = device.createBuffer({
276 label: 'mgpu-readback',
277 size: 4 * Math.max(npts, 2 * nlm),
278 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
279 });
280 const stash = device.createBuffer({
281 label: 'mgpu-state-stash',
282 size: 4 * state.length * 2 * nlm,
283 usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
284 });
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 286 const gpu = new GpuModel({
0af3386Reaction-diffusion on spherical-harmonic surfacesJeremy Magland 287 device, host, initPlan, stepPlan, readback, stash,
288 paramNames, state, view, npts, nlm,
289 });
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 290 if (geometry) gpu.#jhat = geometry.Jhat;
291 return gpu;
294 setParams(params: ModelParams): void {
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 295 const merged = { jhat: this.#jhat, ...params };
296 this.#initPlan.setParams(merged);
297 this.#stepPlan.setParams(merged);
300 /**
301 * Write a host-owned value directly — the spectral state, or one of the input
302 * fields. Lets a test set up an exact initial condition (a single spherical-
303 * harmonic mode, say) instead of going through `init`.
304 */
305 upload(name: string, data: Float32Array): void {
306 this.#host.upload(name, data);
307 }
309 /**
310 * Swap the surface under a running model. The geometry is data, not code —
311 * its shape in the bindings depends only on the grid — so changing it is six
312 * buffer writes and needs no recompile, and the simulation carries straight
313 * on. Only meaningful if the .m took the geometry as an argument.
314 */
315 uploadGeometry(geometry: GeometryBuffers): void {
316 const fields: [string, Float32Array][] = [
317 ['gx', geometry.x], ['gy', geometry.y], ['gz', geometry.z],
318 ['Gx', geometry.X], ['Gy', geometry.Y], ['Gz', geometry.Z],
8144287WIP: added code for Laplace-Beltrami operator evaluation on smooth genus-0 surfaceOwen Melia 319 ['Vtx', geometry.Vtx], ['Vty', geometry.Vty], ['Vtz', geometry.Vtz],
320 ['Vpx', geometry.Vpx], ['Vpy', geometry.Vpy], ['Vpz', geometry.Vpz],
591a4f5Reduce the Laplace-Beltrami matvec to 6 transforms per species per iterationDan Fortunato 321 ['p1', geometry.p1], ['p2', geometry.p2],
322 ['q2', geometry.q2], ['r', geometry.r],
324 for (const [name, data] of fields) {
325 if (this.#host.get(name)) this.#host.upload(name, data);
326 }
e4d6a3bPrecondition with the operator's symbol; project the correction onto the bandDan Fortunato 327 // The new surface's preconditioner scale takes effect on the next
328 // setParams (the session re-applies its params after a swap).
329 this.#jhat = geometry.Jhat;
0ae15cfSeed runs from smooth random fields, and add the blob geometryDan Fortunato 332 /** The wavelength this model's `init` asked `randnfun3` for, or null if it
333 * seeds some other way. The session resolves it and draws the modes. */
334 get randnfun3Lambda(): Randnfun3Lambda | null {
335 return this.#initPlan.randnfun3Lambda;
336 }
338 /**
339 * Upload the seeded initial data and run `init`.
340 *
341 * Both inputs are optional in the sense that a .m uses one or the other:
342 * `modes` is the random field's coefficient table for a model that calls
343 * `randnfun3`, `noise` the plain grid field for one that takes `noise`
344 * directly (the analytic test models inject exact initial conditions that
345 * way). Only what the plan actually bound is uploaded.
346 */
347 async init(noise: Float32Array, modes: Float32Array | null): Promise<void> {
348 if (this.#host.get('noise')) this.#host.upload('noise', noise);
349 // Sized to the wavelength, so this may reallocate and rebind.
350 if (modes) this.#initPlan.uploadRandnfun3Table(this.#host, modes);
351 // Submitted in pieces: a fine seed wavelength makes the mode sum long
352 // enough that one submission would stall the browser's compositor.
353 await this.#initPlan.submitYielding('mgpu-init');
355 }
357 /**
358 * Copy the spectral state aside, so a batch of steps can run — to be timed —
359 * and then be undone with restoreState, leaving the simulation exactly where
360 * it was. Only the state is stashed: the grid view fields keep whatever the
361 * batch last wrote until a subsequent step recomputes them, so step before
362 * reading a view after a restore.
363 */
364 snapshotState(): void {
365 this.#stashedRan = this.#lastRan;
366 this.#copyState('save');
367 }
369 restoreState(): void {
370 this.#copyState('restore');
371 this.#lastRan = this.#stashedRan;
372 }
374 #copyState(dir: 'save' | 'restore'): void {
375 // A restore can land after a rebuild destroyed the buffers mid-await;
376 // there is nothing left to protect, so do not submit into destroyed state.
377 if (this.#destroyed) return;
378 const enc = this.#device.createCommandEncoder({ label: `mgpu-state-${dir}` });
379 let offset = 0;
380 for (const name of this.state) {
381 const slot = this.#host.get(name);
382 if (!slot) throw new Error(`state '${name}' has no host buffer`);
383 const bytes = 4 * slot.count;
384 if (dir === 'save') {
385 enc.copyBufferToBuffer(slot.buffer, 0, this.#stash, offset, bytes);
386 } else {
387 enc.copyBufferToBuffer(this.#stash, offset, slot.buffer, 0, bytes);
388 }
389 offset += bytes;
390 }
391 this.#device.queue.submit([enc.finish()]);
392 }
394 /**
395 * Advance `steps` timesteps. Synchronous — this only records commands and
396 * submits them; nothing is read back and nothing is awaited.
397 */
398 step(steps = 1): void {
399 const enc = this.#device.createCommandEncoder({ label: 'mgpu-step' });
400 this.#stepPlan.encodeSteps(enc, steps);
401 this.#device.queue.submit([enc.finish()]);
402 this.#lastRan = 'step';
403 }
405 /**
406 * The buffer currently holding a named value. Grid fields like `u` are
407 * produced by both functions, into separate buffers (only the spectral state
408 * is shared), so this resolves to whichever function ran most recently —
409 * which is what makes the first frame show the initial state rather than an
410 * unwritten buffer.
411 */
412 #locate(name: string): { buffer: GPUBuffer; count: number } | null {
413 const [first, second] =
414 this.#lastRan === 'init'
415 ? [this.#initPlan, this.#stepPlan]
416 : [this.#stepPlan, this.#initPlan];
417 const buffer = first.buffer(name) ?? second.buffer(name);
418 const count = first.elementCount(name) ?? second.elementCount(name);
419 if (!buffer || count === undefined) return null;
420 return { buffer, count };
421 }
423 /** The GPU buffer a named value would be read from right now — for encoding
424 * further GPU work against it (e.g. a display-grid synthesis of the state)
425 * without a CPU round trip. */
426 valueBuffer(name: string): GPUBuffer | null {
427 return this.#locate(name)?.buffer ?? null;
428 }
430 /** Read a named value back to the CPU. The only await in the whole loop. */
431 async read(name: string): Promise<Float32Array> {
432 const located = this.#locate(name);
433 if (!located) {
434 throw new Error(`read: the model has no value named '${name}'`);
435 }
436 const { buffer, count } = located;
437 const enc = this.#device.createCommandEncoder({ label: `mgpu-read-${name}` });
438 enc.copyBufferToBuffer(buffer, 0, this.#readback, 0, 4 * count);
439 this.#device.queue.submit([enc.finish()]);
440 await this.#readback.mapAsync(GPUMapMode.READ, 0, 4 * count);
441 const out = new Float32Array(this.#readback.getMappedRange(0, 4 * count).slice(0));
442 this.#readback.unmap();
443 return out;
444 }
446 /** What the .m compiled to, for display. */
447 describe(): { init: string[]; step: string[] } {
448 return { init: this.#initPlan.describe(), step: this.#stepPlan.describe() };
449 }
451 destroy(): void {
452 this.#destroyed = true;
453 this.#initPlan.destroy();
454 this.#stepPlan.destroy();
455 this.#host.destroy();
456 this.#readback.destroy();
457 this.#stash.destroy();
458 }
459}