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
5 * the initial state and a `step` function that advances it one timestep. Each
6 * is specialized for the current grid and compiled into a ModelPlan, and both
7 * operate on the same state buffers (see HostBuffers).
8 *
9 * Both functions return the state, in the same order, so their signatures say
10 * exactly what they produce:
11 *
12 * function [p, pm, t] = init(npts)
13 * function [pn, pold, tn] = step(p, pm, t, c, sig, x, y, dt, f, amp, ...)
14 *
15 * The host supplies the things that are setup rather than algorithm: the grid
16 * coordinates, the medium the scene defines, the timestep, and the parameter
17 * values. Each argument is matched to the .m's declared parameter name, so the
18 * file documents its own interface.
19 *
20 * `t` is carried as a grid field rather than a scalar, and that is deliberate.
21 * A batch of timesteps is one replay of a fixed op sequence, so nothing the
22 * host writes between steps can change inside it — a clock uploaded per frame
23 * would stand still for the whole batch. Making the model advance its own time
24 * (`tn = t + dt`) keeps the source term correct however many steps are batched,
25 * at the cost of one extra buffer and one extra kernel per step, which next to
26 * the stencil is nothing.
27 */
28import { HostBuffers, ModelPlan } from './plan.ts';
29import type { StencilPlan } from './stencil.ts';
30import { inFunction, inFunctionAsync, inModel } from './errors.ts';
31import { CompiledModel, type Binding } from './compile.ts';
33export interface ModelParams {
34 [key: string]: number;
35}
37/** The grid a model runs on, and the medium it runs in. */
38export interface GridFields {
39 nx: number;
40 ny: number;
41 /** Grid spacing, the same in x and y. */
42 h: number;
43 /** Coordinates of every grid point, npts each, x fastest. */
44 x: Float32Array;
45 y: Float32Array;
46}
48/** What the scene defines, on the grid. */
49export interface MediumFields {
50 /** Sound speed, npts. */
51 c: Float32Array;
52 /** Absorption rate (the sponge and any absorbing scatterer), npts. */
53 sig: Float32Array;
54}
56export interface GpuModelOptions {
57 device: GPUDevice;
58 stencil: StencilPlan;
59 grid: GridFields;
60 medium: MediumFields;
61 /** Model source (.m text). */
62 source: string;
63 /** Parameter names the .m may take as arguments. */
64 paramNames: string[];
65 /** State field names, in order (e.g. ['p', 'pm', 't']). */
66 state: string[];
67 /** Grid fields one kernel may read, overriding what the device allows.
68 * Only for tests. */
69 operandBudget?: number;
70}
72/** Names the .m may take for the grid coordinates. */
73export const GRID_NAMES = ['x', 'y'] as const;
74/** Names the .m may take for the medium the scene defines. */
75export const MEDIUM_NAMES = ['c', 'sig'] as const;
77export class GpuModel {
78 readonly paramNames: string[];
79 readonly state: string[];
80 readonly npts: number;
82 #device: GPUDevice;
83 #host: HostBuffers;
84 #initPlan: ModelPlan;
85 #stepPlan: ModelPlan;
86 /** Timestep, host-owned: it follows from the medium and the grid (a CFL
87 * condition), not from anything the user types, and it is folded into every
88 * setParams so a .m that takes `dt` is never left with the zero a missing
89 * parameter would default to. */
90 #dt = 0;
91 #readback: GPUBuffer;
92 /** Which function wrote the state most recently; see `read`. */
93 #lastRan: 'init' | 'step' = 'init';
95 private constructor(init: {
96 device: GPUDevice;
97 host: HostBuffers;
98 initPlan: ModelPlan;
99 stepPlan: ModelPlan;
100 readback: GPUBuffer;
101 paramNames: string[];
102 state: string[];
103 npts: number;
104 }) {
105 this.#device = init.device;
106 this.#host = init.host;
107 this.#initPlan = init.initPlan;
108 this.#stepPlan = init.stepPlan;
109 this.#readback = init.readback;
110 this.paramNames = init.paramNames;
111 this.state = init.state;
112 this.npts = init.npts;
113 }
115 static async create(opts: GpuModelOptions): Promise<GpuModel> {
116 const { device, stencil, grid, medium, source, paramNames, state } = opts;
117 const npts = grid.nx * grid.ny;
119 // What the .m may ask for by name. The grid geometry is exact, so a
120 // constructor reading it (`zeros(npts, 1)`) keeps a static shape.
121 const bindings: Record<string, Binding> = {
122 npts: { kind: 'const', value: npts },
123 nx: { kind: 'const', value: grid.nx },
124 ny: { kind: 'const', value: grid.ny },
125 h: { kind: 'const', value: grid.h },
126 dt: { kind: 'param' },
127 };
128 for (const g of GRID_NAMES) bindings[g] = { kind: 'tensor', shape: [npts, 1] };
129 for (const m of MEDIUM_NAMES) bindings[m] = { kind: 'tensor', shape: [npts, 1] };
130 for (const s of state) bindings[s] = { kind: 'tensor', shape: [npts, 1] };
131 for (const p of paramNames) bindings[p] = { kind: 'param' };
133 // Parsing belongs to the file, not to either function.
134 const compiled = inModel(() => new CompiledModel(source, bindings, { npts }));
135 const nargout = state.length;
136 const initFn = inFunction('init', () => compiled.specialize('init', nargout));
137 const stepFn = inFunction('step', () => compiled.specialize('step', nargout));
138 compiled.finish();
140 // Both functions return the state, in order, and both feed it back into
141 // the shared buffers.
142 const feedback = [...state];
144 const host = new HostBuffers(device);
145 // The host owns the state and the inputs it uploads, whether or not a
146 // given function happens to take them as arguments — `init` does not read
147 // `p`, but it writes it, and `step` reads it back.
148 for (const s of state) host.ensure(s, npts);
149 for (const g of GRID_NAMES) host.ensure(g, npts);
150 for (const m of MEDIUM_NAMES) host.ensure(m, npts);
152 const initPlan = await inFunctionAsync('init', () =>
153 ModelPlan.create(device, stencil, { fn: initFn, feedback }, host, opts.operandBudget),
154 );
155 const stepPlan = await inFunctionAsync('step', () =>
156 ModelPlan.create(device, stencil, { fn: stepFn, feedback }, host, opts.operandBudget),
157 );
159 host.upload('x', grid.x);
160 host.upload('y', grid.y);
161 host.upload('c', medium.c);
162 host.upload('sig', medium.sig);
164 const readback = device.createBuffer({
165 label: 'mgpu-readback',
166 size: 4 * npts,
167 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
168 });
170 return new GpuModel({
171 device, host, initPlan, stepPlan, readback, paramNames, state, npts,
172 });
173 }
175 /** The timestep in force. Host-owned; see `#dt`. */
176 get dt(): number {
177 return this.#dt;
178 }
180 setDt(dt: number): void {
181 this.#dt = dt;
182 }
184 setParams(params: ModelParams): void {
185 const merged = { dt: this.#dt, ...params };
186 this.#initPlan.setParams(merged);
187 this.#stepPlan.setParams(merged);
188 }
190 /**
191 * Swap the medium under a running model. It is data, not code — its shape in
192 * the bindings depends only on the grid — so changing the scene is two
193 * buffer writes and needs no recompile.
194 */
195 uploadMedium(medium: MediumFields): void {
196 this.#host.upload('c', medium.c);
197 this.#host.upload('sig', medium.sig);
198 }
200 /** Write a host-owned value directly. Lets a test set up an exact initial
201 * condition instead of going through `init`. */
202 upload(name: string, data: Float32Array): void {
203 this.#host.upload(name, data);
204 }
206 /** Run `init`, replacing the state. */
207 init(): void {
208 const enc = this.#device.createCommandEncoder({ label: 'mgpu-init' });
209 this.#initPlan.encodeSteps(enc, 1);
210 this.#device.queue.submit([enc.finish()]);
211 this.#lastRan = 'init';
212 }
214 /**
215 * Advance `steps` timesteps. Synchronous — this only records commands and
216 * submits them; nothing is read back and nothing is awaited.
217 *
218 * `after` is recorded once per step, so anything that must see every
219 * timestep (the microphone) rides along in the same submission.
220 */
221 step(steps = 1, after?: (encoder: GPUCommandEncoder) => void): void {
222 const enc = this.#device.createCommandEncoder({ label: 'mgpu-step' });
223 this.#stepPlan.encodeSteps(enc, steps, after);
224 this.#device.queue.submit([enc.finish()]);
225 this.#lastRan = 'step';
226 }
228 /**
229 * The buffer currently holding a named value. A field the .m computes is
230 * produced by both functions, into separate buffers (only the state is
231 * shared), so this resolves to whichever function ran most recently — which
232 * is what makes the first frame show the initial state rather than an
233 * unwritten buffer.
234 */
235 #locate(name: string): { buffer: GPUBuffer; count: number } | null {
236 const [first, second] =
237 this.#lastRan === 'init'
238 ? [this.#initPlan, this.#stepPlan]
239 : [this.#stepPlan, this.#initPlan];
240 const buffer = first.buffer(name) ?? second.buffer(name);
241 const count = first.elementCount(name) ?? second.elementCount(name);
242 if (!buffer || count === undefined) return null;
243 return { buffer, count };
244 }
246 /** The GPU buffer a named value would be read from right now. */
247 valueBuffer(name: string): GPUBuffer | null {
248 return this.#locate(name)?.buffer ?? null;
249 }
251 /**
252 * The buffer a host-owned field lives in — the state between calls, or an
253 * input like the sound speed.
254 *
255 * This is what the renderer binds, and it must be this rather than
256 * `valueBuffer`: a bind group is built once and holds a particular buffer,
257 * while `init` and `step` write their outputs into buffers of their own and
258 * only agree here, where their feedback copies land. Binding either
259 * function's private buffer would draw a stale field for half the run.
260 */
261 stateBuffer(name: string): GPUBuffer | null {
262 return this.#host.get(name)?.buffer ?? null;
263 }
265 /** Read a named value back to the CPU. The only await in the whole loop. */
266 async read(name: string): Promise<Float32Array> {
267 const located = this.#locate(name);
268 if (!located) throw new Error(`read: the model has no value named '${name}'`);
269 const { buffer, count } = located;
270 const enc = this.#device.createCommandEncoder({ label: `mgpu-read-${name}` });
271 enc.copyBufferToBuffer(buffer, 0, this.#readback, 0, 4 * count);
272 this.#device.queue.submit([enc.finish()]);
273 await this.#readback.mapAsync(GPUMapMode.READ, 0, 4 * count);
274 const out = new Float32Array(this.#readback.getMappedRange(0, 4 * count).slice(0));
275 this.#readback.unmap();
276 return out;
277 }
279 /** What the .m compiled to, for display. */
280 describe(): { init: string[]; step: string[] } {
281 return { init: this.#initPlan.describe(), step: this.#stepPlan.describe() };
282 }
284 destroy(): void {
285 this.#initPlan.destroy();
286 this.#stepPlan.destroy();
287 this.#host.destroy();
288 this.#readback.destroy();
289 }
290}