1/**
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';
21import type { DerivPlan } from '../sht/deriv.ts';
22import { lmIndex, type ShtConfig } from '../sht/layout.ts';
23import { HostBuffers, ModelPlan } from './plan.ts';
24import { inFunction, inFunctionAsync, inModel } from './errors.ts';
25import { CompiledModel, type Binding } from './compile.ts';
26import { DEFAULT_SOLVER, modelLibs, solveShim, type LibFile } from './libs.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;
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;
54 /**
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 /**
61 * The shared .m files compiled alongside the model. Defaults to the shipped
62 * operator and solvers plus a `solve(...)` shim forwarding to richardson;
63 * the session passes its own list to honor the app's solver selection and
64 * any lib edits.
65 */
66 libs?: LibFile[];
67}
69/** Host-supplied surface fields, in the layout the .m sees them. */
70export interface GeometryBuffers {
71 /** Grid coordinates, npts each. */
72 x: Float32Array;
73 y: Float32Array;
74 z: Float32Array;
75 /** Their spherical-harmonic coefficients, 2 x nlm each. */
76 X: Float32Array;
77 Y: Float32Array;
78 Z: Float32Array;
79 /** Inverse metric quantities (src/geom/metric.ts), grid space, npts each. */
80 Vtx: Float32Array;
81 Vty: Float32Array;
82 Vtz: Float32Array;
83 Vpx: Float32Array;
84 Vpy: Float32Array;
85 Vpz: Float32Array;
86}
88/** Names the .m may take for the grid coordinates and for their coefficients. */
89export const GEOMETRY_GRID_NAMES = ['gx', 'gy', 'gz'] as const;
90export const GEOMETRY_SPECTRAL_NAMES = ['Gx', 'Gy', 'Gz'] as const;
91/** Names the .m may take for the inverse metric quantities. */
92export const METRIC_GRID_NAMES = ['Vtx', 'Vty', 'Vtz', 'Vpx', 'Vpy', 'Vpz'] as const;
94/** Laplace-Beltrami eigenvalues l(l+1), duplicated across re/im so the array
95 * matches the 2 x nlm spectral layout element for element. */
96export function eigenvalues(cfg: ShtConfig, nlm: number): Float32Array {
97 const lam = new Float32Array(2 * nlm);
98 for (let m = 0; m <= cfg.mmax; m++) {
99 for (let l = m; l <= cfg.lmax; l++) {
100 const i = lmIndex(cfg.lmax, l, m);
101 lam[2 * i] = l * (l + 1);
102 lam[2 * i + 1] = l * (l + 1);
103 }
104 }
105 return lam;
106}
108/**
109 * 1 where l < lmax-2, else 0, duplicated across re/im like `lam`. The
110 * theta/phi derivative recurrences (src/sht/derivCoeffs.ts) cannot exactly
111 * represent a derivative at the top two degrees of the band limit, so the
112 * surface Laplace-Beltrami correction filters them out wherever it
113 * re-differentiates a field (evolving_surface/notes/algos.tex Sec 6,
114 * "Miscellaneous implementation details").
115 */
116export function filterMask(cfg: ShtConfig, nlm: number): Float32Array {
117 const filt = new Float32Array(2 * nlm);
118 for (let m = 0; m <= cfg.mmax; m++) {
119 for (let l = m; l <= cfg.lmax; l++) {
120 const i = lmIndex(cfg.lmax, l, m);
121 const keep = l < cfg.lmax - 2 ? 1 : 0;
122 filt[2 * i] = keep;
123 filt[2 * i + 1] = keep;
124 }
125 }
126 return filt;
127}
129/**
130 * Inner-product weights for the half-spectrum layout: 1 at m = 0, 2 at
131 * m > 0, duplicated across re/im like `lam`. The 2 x nlm state stores only
132 * m >= 0 of a real field (the m < 0 coefficients are conjugates), so the
133 * true L2 inner product on the sphere is sum(wlm .* a .* b) — which is what
134 * a Krylov solver's `dot` calls should compute if its scalars are to mean
135 * what they mean in the analysis.
136 */
137export function weightMask(cfg: ShtConfig, nlm: number): Float32Array {
138 const wlm = new Float32Array(2 * nlm);
139 for (let m = 0; m <= cfg.mmax; m++) {
140 for (let l = m; l <= cfg.lmax; l++) {
141 const i = lmIndex(cfg.lmax, l, m);
142 const w = m === 0 ? 1 : 2;
143 wlm[2 * i] = w;
144 wlm[2 * i + 1] = w;
145 }
146 }
147 return wlm;
148}
150export class GpuModel {
151 readonly paramNames: string[];
152 readonly state: string[];
153 readonly view: string[];
154 readonly npts: number;
155 readonly nlm: number;
157 #device: GPUDevice;
158 #host: HostBuffers;
159 #initPlan: ModelPlan;
160 #stepPlan: ModelPlan;
161 #readback: GPUBuffer;
162 /** Scratch holding a copy of the whole spectral state; see snapshotState. */
163 #stash: GPUBuffer;
164 /** Which function wrote the state most recently; see `read`. */
165 #lastRan: 'init' | 'step' = 'init';
166 #stashedRan: 'init' | 'step' = 'init';
167 #destroyed = false;
169 private constructor(init: {
170 device: GPUDevice;
171 host: HostBuffers;
172 initPlan: ModelPlan;
173 stepPlan: ModelPlan;
174 readback: GPUBuffer;
175 stash: GPUBuffer;
176 paramNames: string[];
177 state: string[];
178 view: string[];
179 npts: number;
180 nlm: number;
181 }) {
182 this.#device = init.device;
183 this.#host = init.host;
184 this.#initPlan = init.initPlan;
185 this.#stepPlan = init.stepPlan;
186 this.#readback = init.readback;
187 this.#stash = init.stash;
188 this.paramNames = init.paramNames;
189 this.state = init.state;
190 this.view = init.view;
191 this.npts = init.npts;
192 this.nlm = init.nlm;
193 }
195 static async create(opts: GpuModelOptions): Promise<GpuModel> {
196 const { device, sht, cfg, source, paramNames, state, view, geometry, deriv } = opts;
197 const npts = cfg.nlat * cfg.nphi;
198 const nlm = sht.nlm;
199 const niter = opts.niter ?? 0;
201 // What the .m may ask for by parameter name. Spectral state, the
202 // eigenvalues and the top-mode filter are 2 x nlm; the seeded
203 // perturbation is a grid field.
204 const bindings: Record<string, Binding> = {
205 lam: { kind: 'tensor', shape: [2, nlm] },
206 filt: { kind: 'tensor', shape: [2, nlm] },
207 wlm: { kind: 'tensor', shape: [2, nlm] },
208 noise: { kind: 'tensor', shape: [npts, 1] },
209 npts: { kind: 'const', value: npts },
210 nlm: { kind: 'const', value: nlm },
211 niter: { kind: 'const', value: niter },
212 };
213 if (geometry) {
214 for (const g of GEOMETRY_GRID_NAMES) bindings[g] = { kind: 'tensor', shape: [npts, 1] };
215 for (const g of GEOMETRY_SPECTRAL_NAMES) bindings[g] = { kind: 'tensor', shape: [2, nlm] };
216 for (const g of METRIC_GRID_NAMES) bindings[g] = { kind: 'tensor', shape: [npts, 1] };
217 }
218 for (const s of state) bindings[s] = { kind: 'tensor', shape: [2, nlm] };
219 for (const p of paramNames) bindings[p] = { kind: 'param' };
221 // Parsing belongs to the file, not to either function.
222 const libs = opts.libs ?? [...modelLibs, solveShim(DEFAULT_SOLVER)];
223 const compiled = inModel(
224 () => new CompiledModel(source, bindings, { npts, nlm }, 'model.m', libs),
225 );
226 // Both functions return the new state first, then the rendered grid fields.
227 const nargout = state.length + view.length;
228 const initFn = inFunction('init', () => compiled.specialize('init', nargout));
229 const stepFn = inFunction('step', () => compiled.specialize('step', nargout));
230 compiled.finish();
232 // Only the state outputs feed back into the argument buffers; the grid
233 // fields are read for display and then overwritten next call.
234 const feedback = [...state, ...view.map(() => null)];
236 const host = new HostBuffers(device);
237 // The host owns the state and the inputs it uploads, whether or not a given
238 // function happens to take them as arguments — `init` does not read `U`, but
239 // it writes it, and `step` reads it back.
240 for (const s of state) host.ensure(s, 2 * nlm);
241 host.ensure('lam', 2 * nlm);
242 host.ensure('filt', 2 * nlm);
243 host.ensure('wlm', 2 * nlm);
244 host.ensure('noise', npts);
245 if (geometry) {
246 for (const g of GEOMETRY_GRID_NAMES) host.ensure(g, npts);
247 for (const g of GEOMETRY_SPECTRAL_NAMES) host.ensure(g, 2 * nlm);
248 for (const g of METRIC_GRID_NAMES) host.ensure(g, npts);
249 }
251 const initPlan = await inFunctionAsync('init', () =>
252 ModelPlan.create(device, sht, { fn: initFn, feedback }, host, deriv),
253 );
254 const stepPlan = await inFunctionAsync('step', () =>
255 ModelPlan.create(device, sht, { fn: stepFn, feedback }, host, deriv),
256 );
258 host.upload('lam', eigenvalues(cfg, nlm));
259 host.upload('filt', filterMask(cfg, nlm));
260 host.upload('wlm', weightMask(cfg, nlm));
261 if (geometry) {
262 host.upload('gx', geometry.x);
263 host.upload('gy', geometry.y);
264 host.upload('gz', geometry.z);
265 host.upload('Gx', geometry.X);
266 host.upload('Gy', geometry.Y);
267 host.upload('Gz', geometry.Z);
268 host.upload('Vtx', geometry.Vtx);
269 host.upload('Vty', geometry.Vty);
270 host.upload('Vtz', geometry.Vtz);
271 host.upload('Vpx', geometry.Vpx);
272 host.upload('Vpy', geometry.Vpy);
273 host.upload('Vpz', geometry.Vpz);
274 }
276 const readback = device.createBuffer({
277 label: 'mgpu-readback',
278 size: 4 * Math.max(npts, 2 * nlm),
279 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
280 });
281 const stash = device.createBuffer({
282 label: 'mgpu-state-stash',
283 size: 4 * state.length * 2 * nlm,
284 usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
285 });
287 return new GpuModel({
288 device, host, initPlan, stepPlan, readback, stash,
289 paramNames, state, view, npts, nlm,
290 });
291 }
293 setParams(params: ModelParams): void {
294 this.#initPlan.setParams(params);
295 this.#stepPlan.setParams(params);
296 }
298 /**
299 * Write a host-owned value directly — the spectral state, or one of the input
300 * fields. Lets a test set up an exact initial condition (a single spherical-
301 * harmonic mode, say) instead of going through `init`.
302 */
303 upload(name: string, data: Float32Array): void {
304 this.#host.upload(name, data);
305 }
307 /**
308 * Swap the surface under a running model. The geometry is data, not code —
309 * its shape in the bindings depends only on the grid — so changing it is six
310 * buffer writes and needs no recompile, and the simulation carries straight
311 * on. Only meaningful if the .m took the geometry as an argument.
312 */
313 uploadGeometry(geometry: GeometryBuffers): void {
314 const fields: [string, Float32Array][] = [
315 ['gx', geometry.x], ['gy', geometry.y], ['gz', geometry.z],
316 ['Gx', geometry.X], ['Gy', geometry.Y], ['Gz', geometry.Z],
317 ['Vtx', geometry.Vtx], ['Vty', geometry.Vty], ['Vtz', geometry.Vtz],
318 ['Vpx', geometry.Vpx], ['Vpy', geometry.Vpy], ['Vpz', geometry.Vpz],
319 ];
320 for (const [name, data] of fields) {
321 if (this.#host.get(name)) this.#host.upload(name, data);
322 }
323 }
325 /** Upload the seeded perturbation and run `init`. */
326 init(noise: Float32Array): void {
327 this.#host.upload('noise', noise);
328 const enc = this.#device.createCommandEncoder({ label: 'mgpu-init' });
329 this.#initPlan.encodeSteps(enc, 1);
330 this.#device.queue.submit([enc.finish()]);
331 this.#lastRan = 'init';
332 }
334 /**
335 * Copy the spectral state aside, so a batch of steps can run — to be timed —
336 * and then be undone with restoreState, leaving the simulation exactly where
337 * it was. Only the state is stashed: the grid view fields keep whatever the
338 * batch last wrote until a subsequent step recomputes them, so step before
339 * reading a view after a restore.
340 */
341 snapshotState(): void {
342 this.#stashedRan = this.#lastRan;
343 this.#copyState('save');
344 }
346 restoreState(): void {
347 this.#copyState('restore');
348 this.#lastRan = this.#stashedRan;
349 }
351 #copyState(dir: 'save' | 'restore'): void {
352 // A restore can land after a rebuild destroyed the buffers mid-await;
353 // there is nothing left to protect, so do not submit into destroyed state.
354 if (this.#destroyed) return;
355 const enc = this.#device.createCommandEncoder({ label: `mgpu-state-${dir}` });
356 let offset = 0;
357 for (const name of this.state) {
358 const slot = this.#host.get(name);
359 if (!slot) throw new Error(`state '${name}' has no host buffer`);
360 const bytes = 4 * slot.count;
361 if (dir === 'save') {
362 enc.copyBufferToBuffer(slot.buffer, 0, this.#stash, offset, bytes);
363 } else {
364 enc.copyBufferToBuffer(this.#stash, offset, slot.buffer, 0, bytes);
365 }
366 offset += bytes;
367 }
368 this.#device.queue.submit([enc.finish()]);
369 }
371 /**
372 * Advance `steps` timesteps. Synchronous — this only records commands and
373 * submits them; nothing is read back and nothing is awaited.
374 */
375 step(steps = 1): void {
376 const enc = this.#device.createCommandEncoder({ label: 'mgpu-step' });
377 this.#stepPlan.encodeSteps(enc, steps);
378 this.#device.queue.submit([enc.finish()]);
379 this.#lastRan = 'step';
380 }
382 /**
383 * The buffer currently holding a named value. Grid fields like `u` are
384 * produced by both functions, into separate buffers (only the spectral state
385 * is shared), so this resolves to whichever function ran most recently —
386 * which is what makes the first frame show the initial state rather than an
387 * unwritten buffer.
388 */
389 #locate(name: string): { buffer: GPUBuffer; count: number } | null {
390 const [first, second] =
391 this.#lastRan === 'init'
392 ? [this.#initPlan, this.#stepPlan]
393 : [this.#stepPlan, this.#initPlan];
394 const buffer = first.buffer(name) ?? second.buffer(name);
395 const count = first.elementCount(name) ?? second.elementCount(name);
396 if (!buffer || count === undefined) return null;
397 return { buffer, count };
398 }
400 /** The GPU buffer a named value would be read from right now — for encoding
401 * further GPU work against it (e.g. a display-grid synthesis of the state)
402 * without a CPU round trip. */
403 valueBuffer(name: string): GPUBuffer | null {
404 return this.#locate(name)?.buffer ?? null;
405 }
407 /** Read a named value back to the CPU. The only await in the whole loop. */
408 async read(name: string): Promise<Float32Array> {
409 const located = this.#locate(name);
410 if (!located) {
411 throw new Error(`read: the model has no value named '${name}'`);
412 }
413 const { buffer, count } = located;
414 const enc = this.#device.createCommandEncoder({ label: `mgpu-read-${name}` });
415 enc.copyBufferToBuffer(buffer, 0, this.#readback, 0, 4 * count);
416 this.#device.queue.submit([enc.finish()]);
417 await this.#readback.mapAsync(GPUMapMode.READ, 0, 4 * count);
418 const out = new Float32Array(this.#readback.getMappedRange(0, 4 * count).slice(0));
419 this.#readback.unmap();
420 return out;
421 }
423 /** What the .m compiled to, for display. */
424 describe(): { init: string[]; step: string[] } {
425 return { init: this.#initPlan.describe(), step: this.#stepPlan.describe() };
426 }
428 destroy(): void {
429 this.#destroyed = true;
430 this.#initPlan.destroy();
431 this.#stepPlan.destroy();
432 this.#host.destroy();
433 this.#readback.destroy();
434 this.#stash.destroy();
435 }
436}