2 * The .m model, compiled and running on the GPU.
3 *
4 * The model file is ordinary MATLAB: it defines an `init` function that builds
5 * the initial state — each object shaped into its excitation, silent air — and
6 * a `step` function that advances everything one timestep. Each is specialized
7 * for the current grids and compiled into a ModelPlan, and both operate on the
8 * same state buffers (see HostBuffers).
9 *
10 * Both functions return the state, in the same order, so their signatures say
11 * exactly what they produce. The state spans several grids — a pair of history
12 * fields per vibrating object, plus the air's pressure pair — and the compiled
13 * statements mix freely: each line is one kernel over its own field's size,
14 * and the external ops (src/mgpu/ops.ts) are where a value crosses from one
15 * grid to another.
16 *
17 * Where dulcimer hardcoded its two grids by name, everything here is data the
18 * caller supplies: `consts` are the exact scalars a .m may read (grid counts,
19 * spacings, lengths — exact so array constructors keep static shapes),
20 * `fields` are the host-uploaded arrays (coordinates, pin masks, the medium
21 * and the coupling profiles), `state` names the fields the .m advances and how
22 * big each is, and `paramNames` the tunable scalars. The .m names the
23 * arguments it wants and the host matches each against this offer; a name
24 * nobody provides is a compile error with a position, not a silently undefined
25 * variable.
26 *
27 * There is no clock here, and that is not an oversight: every excitation is
28 * an initial condition, not a driven source, so nothing in the model needs to
29 * know what time it is.
30 */
31import { HostBuffers, ModelPlan } from './plan.ts';
32import type { OpPlan } from './ops.ts';
33import { inFunction, inFunctionAsync, inModel } from './errors.ts';
34import { CompiledModel, type Binding } from './compile.ts';
36export interface ModelParams {
37 [key: string]: number;
38}
40/** One state field the .m advances, and its element count. */
41export interface StateField {
42 name: string;
43 count: number;
44}
46export interface GpuModelOptions {
47 device: GPUDevice;
48 ops: OpPlan;
49 /** Exact scalars the .m may take as arguments (grid counts, spacings). */
50 consts: Record<string, number>;
51 /** Host-uploaded arrays the .m may take as arguments (coordinates, masks,
52 * the medium, the coupling profiles), with their initial contents. */
53 fields: Record<string, Float32Array>;
54 /** Model source (.m text). */
55 source: string;
56 /** Tunable scalar names the .m may take as arguments. */
57 paramNames: string[];
58 /** State fields the .m advances, in the order its functions return them. */
59 state: StateField[];
60 /** Grid fields one kernel may read, overriding what the device allows.
61 * Only for tests. */
62 operandBudget?: number;
63}
65export class GpuModel {
66 readonly paramNames: string[];
67 readonly state: StateField[];
69 #device: GPUDevice;
70 #host: HostBuffers;
71 #initPlan: ModelPlan;
72 #stepPlan: ModelPlan;
73 /** Names of the uploadable fields, for uploadFields' checks. */
74 #fieldNames: Set<string>;
75 /** Timestep, host-owned: it follows from the grid and the medium (a CFL
76 * condition), not from anything the user types, and it is folded into every
77 * setParams so the .m's `dt` is never left with the zero a missing
78 * parameter would default to. */
79 #dt = 0;
80 #readback: GPUBuffer;
81 /** Which function wrote the state most recently; see `read`. */
82 #lastRan: 'init' | 'step' = 'init';
84 private constructor(init: {
85 device: GPUDevice;
86 host: HostBuffers;
87 initPlan: ModelPlan;
88 stepPlan: ModelPlan;
89 readback: GPUBuffer;
90 paramNames: string[];
91 state: StateField[];
92 fieldNames: Set<string>;
93 }) {
94 this.#device = init.device;
95 this.#host = init.host;
96 this.#initPlan = init.initPlan;
97 this.#stepPlan = init.stepPlan;
98 this.#readback = init.readback;
99 this.paramNames = init.paramNames;
100 this.state = init.state;
101 this.#fieldNames = init.fieldNames;
102 }
104 static async create(opts: GpuModelOptions): Promise<GpuModel> {
105 const { device, ops, consts, fields, source, paramNames, state } = opts;
107 // What the .m may ask for by name. The grid geometry is exact, so a
108 // constructor reading it (`zeros(npts, 1)`) keeps a static shape.
109 const bindings: Record<string, Binding> = { dt: { kind: 'param' } };
110 for (const [name, value] of Object.entries(consts)) {
111 bindings[name] = { kind: 'const', value };
112 }
113 for (const [name, data] of Object.entries(fields)) {
114 bindings[name] = { kind: 'tensor', shape: [data.length, 1] };
115 }
116 for (const s of state) bindings[s.name] = { kind: 'tensor', shape: [s.count, 1] };
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, ops.instances));
121 const nargout = state.length;
122 const initFn = inFunction('init', () => compiled.specialize('init', nargout));
123 const stepFn = inFunction('step', () => compiled.specialize('step', nargout));
124 compiled.finish();
126 // Both functions return the state, in order, and both feed it back into
127 // the shared buffers.
128 const feedback = state.map((s) => s.name);
130 const host = new HostBuffers(device);
131 // The host owns the state and the inputs it uploads, whether or not a
132 // given function happens to take them as arguments.
133 for (const s of state) host.ensure(s.name, s.count);
134 for (const [name, data] of Object.entries(fields)) host.ensure(name, data.length);
136 const initPlan = await inFunctionAsync('init', () =>
137 ModelPlan.create(device, ops, { fn: initFn, feedback }, host, opts.operandBudget),
138 );
139 const stepPlan = await inFunctionAsync('step', () =>
140 ModelPlan.create(device, ops, { fn: stepFn, feedback }, host, opts.operandBudget),
141 );
143 for (const [name, data] of Object.entries(fields)) host.upload(name, data);
145 const maxCount = Math.max(...state.map((s) => s.count), ...Object.values(fields).map((f) => f.length));
146 const readback = device.createBuffer({
147 label: 'mgpu-readback',
148 size: 4 * maxCount,
149 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
150 });
152 return new GpuModel({
153 device, host, initPlan, stepPlan, readback, paramNames, state,
154 fieldNames: new Set(Object.keys(fields)),
155 });
156 }
158 /** The timestep in force. Host-owned; see `#dt`. */
159 get dt(): number {
160 return this.#dt;
161 }
163 setDt(dt: number): void {
164 this.#dt = dt;
165 }
167 setParams(params: ModelParams): void {
168 const merged = { dt: this.#dt, ...params };
169 this.#initPlan.setParams(merged);
170 this.#stepPlan.setParams(merged);
171 }
173 /**
174 * Swap host-uploaded fields under a running model — the medium and the
175 * coupling profiles after a scene edit. They are data, not code — their
176 * shapes in the bindings depend only on the grids — so changing the body's
177 * geometry or moving an object is a few buffer writes and no recompile.
178 */
179 uploadFields(fields: Record<string, Float32Array>): void {
180 for (const [name, data] of Object.entries(fields)) {
181 if (!this.#fieldNames.has(name)) {
182 throw new Error(`uploadFields: '${name}' is not a host field of this model`);
183 }
184 this.#host.upload(name, data);
185 }
186 }
188 /** Write a host-owned value directly. Lets a test set up an exact initial
189 * condition instead of going through `init`. */
190 upload(name: string, data: Float32Array): void {
191 this.#host.upload(name, data);
192 }
194 /** Run `init`, replacing the state. */
195 init(): void {
196 const enc = this.#device.createCommandEncoder({ label: 'mgpu-init' });
197 this.#initPlan.encodeSteps(enc, 1);
198 this.#device.queue.submit([enc.finish()]);
199 this.#lastRan = 'init';
200 }
202 /**
203 * Advance `steps` timesteps. Synchronous — this only records commands and
204 * submits them; nothing is read back and nothing is awaited.
205 *
206 * `after` is recorded once per step, so anything that must see every
207 * timestep (the microphone) rides along in the same submission.
208 */
209 step(steps = 1, after?: (encoder: GPUCommandEncoder) => void): void {
210 const enc = this.#device.createCommandEncoder({ label: 'mgpu-step' });
211 this.#stepPlan.encodeSteps(enc, steps, after);
212 this.#device.queue.submit([enc.finish()]);
213 this.#lastRan = 'step';
214 }
216 /**
217 * The buffer currently holding a named value. A field the .m computes is
218 * produced by both functions, into separate buffers (only the state is
219 * shared), so this resolves to whichever function ran most recently — which
220 * is what makes the first frame show the initial state rather than an
221 * unwritten buffer.
222 */
223 #locate(name: string): { buffer: GPUBuffer; count: number } | null {
224 const [first, second] =
225 this.#lastRan === 'init'
226 ? [this.#initPlan, this.#stepPlan]
227 : [this.#stepPlan, this.#initPlan];
228 const buffer = first.buffer(name) ?? second.buffer(name);
229 const count = first.elementCount(name) ?? second.elementCount(name);
230 if (!buffer || count === undefined) return null;
231 return { buffer, count };
232 }
234 /** The GPU buffer a named value would be read from right now. */
235 valueBuffer(name: string): GPUBuffer | null {
236 return this.#locate(name)?.buffer ?? null;
237 }
239 /**
240 * The buffer a host-owned field lives in — the state between calls, or an
241 * input like the wall mask.
242 *
243 * This is what the renderer binds, and it must be this rather than
244 * `valueBuffer`: a bind group is built once and holds a particular buffer,
245 * while `init` and `step` write their outputs into buffers of their own and
246 * only agree here, where their feedback copies land. Binding either
247 * function's private buffer would draw a stale field for half the run.
248 */
249 stateBuffer(name: string): GPUBuffer | null {
250 return this.#host.get(name)?.buffer ?? null;
251 }
253 /** Read a named value back to the CPU. The only await in the whole loop. */
254 async read(name: string): Promise<Float32Array> {
255 const located = this.#locate(name);
256 if (!located) throw new Error(`read: the model has no value named '${name}'`);
257 const { buffer, count } = located;
258 const enc = this.#device.createCommandEncoder({ label: `mgpu-read-${name}` });
259 enc.copyBufferToBuffer(buffer, 0, this.#readback, 0, 4 * count);
260 this.#device.queue.submit([enc.finish()]);
261 await this.#readback.mapAsync(GPUMapMode.READ, 0, 4 * count);
262 const out = new Float32Array(this.#readback.getMappedRange(0, 4 * count).slice(0));
263 this.#readback.unmap();
264 return out;
265 }
267 /** What the .m compiled to, for display. */
268 describe(): { init: string[]; step: string[] } {
269 return { init: this.#initPlan.describe(), step: this.#stepPlan.describe() };
270 }
272 destroy(): void {
273 this.#initPlan.destroy();
274 this.#stepPlan.destroy();
275 this.#host.destroy();
276 this.#readback.destroy();
277 }
278}