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 { lmIndex, type ShtConfig } from '../sht/layout.ts';
22import { HostBuffers, ModelPlan } from './plan.ts';
23import { inFunction, inFunctionAsync, inModel } from './errors.ts';
24import { CompiledModel, type Binding } from './compile.ts';
26export interface ModelParams {
27 [key: string]: number;
28}
30export interface GpuModelOptions {
31 device: GPUDevice;
32 sht: ShtPlan;
33 cfg: ShtConfig;
34 /** Model source (.m text). */
35 source: string;
36 /** Parameter names the .m may take as arguments. */
37 paramNames: string[];
38 /** Spectral state names, in order (e.g. ['U', 'V']). */
39 state: string[];
40 /** Grid fields to render, in order (e.g. ['u', 'v']). */
41 view: string[];
42}
44/** Laplace-Beltrami eigenvalues l(l+1), duplicated across re/im so the array
45 * matches the 2 x nlm spectral layout element for element. */
46export function eigenvalues(cfg: ShtConfig, nlm: number): Float32Array {
47 const lam = new Float32Array(2 * nlm);
48 for (let m = 0; m <= cfg.mmax; m++) {
49 for (let l = m; l <= cfg.lmax; l++) {
50 const i = lmIndex(cfg.lmax, l, m);
51 lam[2 * i] = l * (l + 1);
52 lam[2 * i + 1] = l * (l + 1);
53 }
54 }
55 return lam;
56}
58export class GpuModel {
59 readonly paramNames: string[];
60 readonly state: string[];
61 readonly view: string[];
62 readonly npts: number;
63 readonly nlm: number;
65 #device: GPUDevice;
66 #host: HostBuffers;
67 #initPlan: ModelPlan;
68 #stepPlan: ModelPlan;
69 #readback: GPUBuffer;
70 /** Scratch holding a copy of the whole spectral state; see snapshotState. */
71 #stash: GPUBuffer;
72 /** Which function wrote the state most recently; see `read`. */
73 #lastRan: 'init' | 'step' = 'init';
74 #stashedRan: 'init' | 'step' = 'init';
75 #destroyed = false;
77 private constructor(init: {
78 device: GPUDevice;
79 host: HostBuffers;
80 initPlan: ModelPlan;
81 stepPlan: ModelPlan;
82 readback: GPUBuffer;
83 stash: GPUBuffer;
84 paramNames: string[];
85 state: string[];
86 view: string[];
87 npts: number;
88 nlm: number;
89 }) {
90 this.#device = init.device;
91 this.#host = init.host;
92 this.#initPlan = init.initPlan;
93 this.#stepPlan = init.stepPlan;
94 this.#readback = init.readback;
95 this.#stash = init.stash;
96 this.paramNames = init.paramNames;
97 this.state = init.state;
98 this.view = init.view;
99 this.npts = init.npts;
100 this.nlm = init.nlm;
101 }
103 static async create(opts: GpuModelOptions): Promise<GpuModel> {
104 const { device, sht, cfg, source, paramNames, state, view } = opts;
105 const npts = cfg.nlat * cfg.nphi;
106 const nlm = sht.nlm;
108 // What the .m may ask for by parameter name. Spectral state and the
109 // eigenvalues are 2 x nlm; the seeded perturbation is a grid field.
110 const bindings: Record<string, Binding> = {
111 lam: { kind: 'tensor', shape: [2, nlm] },
112 noise: { kind: 'tensor', shape: [npts, 1] },
113 npts: { kind: 'const', value: npts },
114 nlm: { kind: 'const', value: nlm },
115 };
116 for (const s of state) bindings[s] = { kind: 'tensor', shape: [2, nlm] };
117 for (const p of paramNames) bindings[p] = { kind: 'param' };
119 // Parsing belongs to the file, not to either function.
120 const compiled = inModel(() => new CompiledModel(source, bindings, { npts, nlm }));
121 // Both functions return the new state first, then the rendered grid fields.
122 const nargout = state.length + view.length;
123 const initFn = inFunction('init', () => compiled.specialize('init', nargout));
124 const stepFn = inFunction('step', () => compiled.specialize('step', nargout));
125 compiled.finish();
127 // Only the state outputs feed back into the argument buffers; the grid
128 // fields are read for display and then overwritten next call.
129 const feedback = [...state, ...view.map(() => null)];
131 const host = new HostBuffers(device);
132 // The host owns the state and the inputs it uploads, whether or not a given
133 // function happens to take them as arguments — `init` does not read `U`, but
134 // it writes it, and `step` reads it back.
135 for (const s of state) host.ensure(s, 2 * nlm);
136 host.ensure('lam', 2 * nlm);
137 host.ensure('noise', npts);
139 const initPlan = await inFunctionAsync('init', () =>
140 ModelPlan.create(device, sht, { fn: initFn, feedback }, host),
141 );
142 const stepPlan = await inFunctionAsync('step', () =>
143 ModelPlan.create(device, sht, { fn: stepFn, feedback }, host),
144 );
146 host.upload('lam', eigenvalues(cfg, nlm));
148 const readback = device.createBuffer({
149 label: 'mgpu-readback',
150 size: 4 * Math.max(npts, 2 * nlm),
151 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
152 });
153 const stash = device.createBuffer({
154 label: 'mgpu-state-stash',
155 size: 4 * state.length * 2 * nlm,
156 usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
157 });
159 return new GpuModel({
160 device, host, initPlan, stepPlan, readback, stash,
161 paramNames, state, view, npts, nlm,
162 });
163 }
165 setParams(params: ModelParams): void {
166 this.#initPlan.setParams(params);
167 this.#stepPlan.setParams(params);
168 }
170 /**
171 * Write a host-owned value directly — the spectral state, or one of the input
172 * fields. Lets a test set up an exact initial condition (a single spherical-
173 * harmonic mode, say) instead of going through `init`.
174 */
175 upload(name: string, data: Float32Array): void {
176 this.#host.upload(name, data);
177 }
179 /** Upload the seeded perturbation and run `init`. */
180 init(noise: Float32Array): void {
181 this.#host.upload('noise', noise);
182 const enc = this.#device.createCommandEncoder({ label: 'mgpu-init' });
183 this.#initPlan.encodeSteps(enc, 1);
184 this.#device.queue.submit([enc.finish()]);
185 this.#lastRan = 'init';
186 }
188 /**
189 * Copy the spectral state aside, so a batch of steps can run — to be timed —
190 * and then be undone with restoreState, leaving the simulation exactly where
191 * it was. Only the state is stashed: the grid view fields keep whatever the
192 * batch last wrote until a subsequent step recomputes them, so step before
193 * reading a view after a restore.
194 */
195 snapshotState(): void {
196 this.#stashedRan = this.#lastRan;
197 this.#copyState('save');
198 }
200 restoreState(): void {
201 this.#copyState('restore');
202 this.#lastRan = this.#stashedRan;
203 }
205 #copyState(dir: 'save' | 'restore'): void {
206 // A restore can land after a rebuild destroyed the buffers mid-await;
207 // there is nothing left to protect, so do not submit into destroyed state.
208 if (this.#destroyed) return;
209 const enc = this.#device.createCommandEncoder({ label: `mgpu-state-${dir}` });
210 let offset = 0;
211 for (const name of this.state) {
212 const slot = this.#host.get(name);
213 if (!slot) throw new Error(`state '${name}' has no host buffer`);
214 const bytes = 4 * slot.count;
215 if (dir === 'save') {
216 enc.copyBufferToBuffer(slot.buffer, 0, this.#stash, offset, bytes);
217 } else {
218 enc.copyBufferToBuffer(this.#stash, offset, slot.buffer, 0, bytes);
219 }
220 offset += bytes;
221 }
222 this.#device.queue.submit([enc.finish()]);
223 }
225 /**
226 * Advance `steps` timesteps. Synchronous — this only records commands and
227 * submits them; nothing is read back and nothing is awaited.
228 */
229 step(steps = 1): void {
230 const enc = this.#device.createCommandEncoder({ label: 'mgpu-step' });
231 this.#stepPlan.encodeSteps(enc, steps);
232 this.#device.queue.submit([enc.finish()]);
233 this.#lastRan = 'step';
234 }
236 /**
237 * The buffer currently holding a named value. Grid fields like `u` are
238 * produced by both functions, into separate buffers (only the spectral state
239 * is shared), so this resolves to whichever function ran most recently —
240 * which is what makes the first frame show the initial state rather than an
241 * unwritten buffer.
242 */
243 #locate(name: string): { buffer: GPUBuffer; count: number } | null {
244 const [first, second] =
245 this.#lastRan === 'init'
246 ? [this.#initPlan, this.#stepPlan]
247 : [this.#stepPlan, this.#initPlan];
248 const buffer = first.buffer(name) ?? second.buffer(name);
249 const count = first.elementCount(name) ?? second.elementCount(name);
250 if (!buffer || count === undefined) return null;
251 return { buffer, count };
252 }
254 /** The GPU buffer a named value would be read from right now — for encoding
255 * further GPU work against it (e.g. a display-grid synthesis of the state)
256 * without a CPU round trip. */
257 valueBuffer(name: string): GPUBuffer | null {
258 return this.#locate(name)?.buffer ?? null;
259 }
261 /** Read a named value back to the CPU. The only await in the whole loop. */
262 async read(name: string): Promise<Float32Array> {
263 const located = this.#locate(name);
264 if (!located) {
265 throw new Error(`read: the model has no value named '${name}'`);
266 }
267 const { buffer, count } = located;
268 const enc = this.#device.createCommandEncoder({ label: `mgpu-read-${name}` });
269 enc.copyBufferToBuffer(buffer, 0, this.#readback, 0, 4 * count);
270 this.#device.queue.submit([enc.finish()]);
271 await this.#readback.mapAsync(GPUMapMode.READ, 0, 4 * count);
272 const out = new Float32Array(this.#readback.getMappedRange(0, 4 * count).slice(0));
273 this.#readback.unmap();
274 return out;
275 }
277 /** What the .m compiled to, for display. */
278 describe(): { init: string[]; step: string[] } {
279 return { init: this.#initPlan.describe(), step: this.#stepPlan.describe() };
280 }
282 destroy(): void {
283 this.#destroyed = true;
284 this.#initPlan.destroy();
285 this.#stepPlan.destroy();
286 this.#host.destroy();
287 this.#readback.destroy();
288 this.#stash.destroy();
289 }
290}