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 /** Which function wrote the state most recently; see `read`. */
71 #lastRan: 'init' | 'step' = 'init';
73 private constructor(init: {
74 device: GPUDevice;
75 host: HostBuffers;
76 initPlan: ModelPlan;
77 stepPlan: ModelPlan;
78 readback: GPUBuffer;
79 paramNames: string[];
80 state: string[];
81 view: string[];
82 npts: number;
83 nlm: number;
84 }) {
85 this.#device = init.device;
86 this.#host = init.host;
87 this.#initPlan = init.initPlan;
88 this.#stepPlan = init.stepPlan;
89 this.#readback = init.readback;
90 this.paramNames = init.paramNames;
91 this.state = init.state;
92 this.view = init.view;
93 this.npts = init.npts;
94 this.nlm = init.nlm;
95 }
97 static async create(opts: GpuModelOptions): Promise<GpuModel> {
98 const { device, sht, cfg, source, paramNames, state, view } = opts;
99 const npts = cfg.nlat * cfg.nphi;
100 const nlm = sht.nlm;
102 // What the .m may ask for by parameter name. Spectral state and the
103 // eigenvalues are 2 x nlm; the seeded perturbation is a grid field.
104 const bindings: Record<string, Binding> = {
105 lam: { kind: 'tensor', shape: [2, nlm] },
106 noise: { kind: 'tensor', shape: [npts, 1] },
107 npts: { kind: 'const', value: npts },
108 nlm: { kind: 'const', value: nlm },
109 };
110 for (const s of state) bindings[s] = { kind: 'tensor', shape: [2, nlm] };
111 for (const p of paramNames) bindings[p] = { kind: 'param' };
113 // Parsing belongs to the file, not to either function.
114 const compiled = inModel(() => new CompiledModel(source, bindings, { npts, nlm }));
115 // Both functions return the new state first, then the rendered grid fields.
116 const nargout = state.length + view.length;
117 const initFn = inFunction('init', () => compiled.specialize('init', nargout));
118 const stepFn = inFunction('step', () => compiled.specialize('step', nargout));
119 compiled.finish();
121 // Only the state outputs feed back into the argument buffers; the grid
122 // fields are read for display and then overwritten next call.
123 const feedback = [...state, ...view.map(() => null)];
125 const host = new HostBuffers(device);
126 // The host owns the state and the inputs it uploads, whether or not a given
127 // function happens to take them as arguments — `init` does not read `U`, but
128 // it writes it, and `step` reads it back.
129 for (const s of state) host.ensure(s, 2 * nlm);
130 host.ensure('lam', 2 * nlm);
131 host.ensure('noise', npts);
133 const initPlan = await inFunctionAsync('init', () =>
134 ModelPlan.create(device, sht, { fn: initFn, feedback }, host),
135 );
136 const stepPlan = await inFunctionAsync('step', () =>
137 ModelPlan.create(device, sht, { fn: stepFn, feedback }, host),
138 );
140 host.upload('lam', eigenvalues(cfg, nlm));
142 const readback = device.createBuffer({
143 label: 'mgpu-readback',
144 size: 4 * Math.max(npts, 2 * nlm),
145 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
146 });
148 return new GpuModel({
149 device, host, initPlan, stepPlan, readback,
150 paramNames, state, view, npts, nlm,
151 });
152 }
154 setParams(params: ModelParams): void {
155 this.#initPlan.setParams(params);
156 this.#stepPlan.setParams(params);
157 }
159 /** Upload the seeded perturbation and run `init`. */
160 init(noise: Float32Array): void {
161 this.#host.upload('noise', noise);
162 const enc = this.#device.createCommandEncoder({ label: 'mgpu-init' });
163 this.#initPlan.encodeSteps(enc, 1);
164 this.#device.queue.submit([enc.finish()]);
165 this.#lastRan = 'init';
166 }
168 /**
169 * Advance `steps` timesteps. Synchronous — this only records commands and
170 * submits them; nothing is read back and nothing is awaited.
171 */
172 step(steps = 1): void {
173 const enc = this.#device.createCommandEncoder({ label: 'mgpu-step' });
174 this.#stepPlan.encodeSteps(enc, steps);
175 this.#device.queue.submit([enc.finish()]);
176 this.#lastRan = 'step';
177 }
179 /**
180 * Read a named value back to the CPU. The only await in the whole loop.
181 *
182 * Grid fields like `u` are produced by both functions, into separate buffers
183 * (only the spectral state is shared), so this reads from whichever ran most
184 * recently — which is what makes the first frame show the initial state
185 * rather than an unwritten buffer.
186 */
187 async read(name: string): Promise<Float32Array> {
188 const [first, second] =
189 this.#lastRan === 'init'
190 ? [this.#initPlan, this.#stepPlan]
191 : [this.#stepPlan, this.#initPlan];
192 const buffer = first.buffer(name) ?? second.buffer(name);
193 const count = first.elementCount(name) ?? second.elementCount(name);
194 if (!buffer || count === undefined) {
195 throw new Error(`read: the model has no value named '${name}'`);
196 }
197 const enc = this.#device.createCommandEncoder({ label: `mgpu-read-${name}` });
198 enc.copyBufferToBuffer(buffer, 0, this.#readback, 0, 4 * count);
199 this.#device.queue.submit([enc.finish()]);
200 await this.#readback.mapAsync(GPUMapMode.READ, 0, 4 * count);
201 const out = new Float32Array(this.#readback.getMappedRange(0, 4 * count).slice(0));
202 this.#readback.unmap();
203 return out;
204 }
206 /** What the .m compiled to, for display. */
207 describe(): { init: string[]; step: string[] } {
208 return { init: this.#initPlan.describe(), step: this.#stepPlan.describe() };
209 }
211 destroy(): void {
212 this.#initPlan.destroy();
213 this.#stepPlan.destroy();
214 this.#host.destroy();
215 this.#readback.destroy();
216 }
217}