1035139A plucked dulcimer string and its box, as two coupled wave equations on WebGPUJeremy Magland 1/**
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 — the plucked string and a silent air field — and a `step`
6 * function that advances both one timestep. Each is specialized for the
7 * current grids and compiled into a ModelPlan, and both operate on the same
8 * state buffers (see HostBuffers).
9 *
10 * Both functions return the state, in the same order, so their signatures say
11 * exactly what they produce:
12 *
13 * function [u, um, p, pm] = init(xs, npts, Ls, pluckpos, amp)
14 * function [un, uold, pn, pold] = step(u, um, p, pm, ...)
15 *
16 * The state spans two grids — `u`, `um` are string fields (ns nodes), `p`,
17 * `pm` air fields (npts cells) — and the compiled statements mix freely:
18 * each line is one kernel over its own field's size, and the external ops
19 * (src/mgpu/ops.ts) are where a value crosses from one grid to the other.
20 *
21 * The host supplies the things that are setup rather than algorithm: the grid
22 * coordinates, the medium and coupling profiles the scene defines, the
23 * timestep, and the parameter values. Each argument is matched to the .m's
24 * declared parameter name, so the file documents its own interface.
25 *
26 * There is no clock here, and that is not an oversight: the pluck is an
27 * initial condition, not a driven source, so nothing in the model needs to
28 * know what time it is.
29 */
30import { HostBuffers, ModelPlan } from './plan.ts';
31import type { OpPlan } from './ops.ts';
32import { inFunction, inFunctionAsync, inModel } from './errors.ts';
33import { CompiledModel, type Binding } from './compile.ts';
34import type { AirGrid, StringGrid } from '../grid.ts';
36export interface ModelParams {
37 [key: string]: number;
38}
40/** What the scene defines, on the air grid: npts values each. */
41export interface MediumFields {
42 /** Sound speed, m/s. */
43 c: Float32Array;
44 /** Absorption rate, 1/s (the sponge and any wall absorption). */
45 sig: Float32Array;
46 /** 1 in air, 0 in the body's solid shell. What `lapw` masks by. */
47 wall: Float32Array;
48 /** Where the string radiates directly: a tube around the string line. */
49 lineprof: Float32Array;
50 /** Where the bridge force drives the air: a patch above the top plate. */
51 boardprof: Float32Array;
52}
54/** One state field the .m advances, and which grid it lives on. */
55export interface StateField {
56 name: string;
57 grid: 'air' | 'string';
58}
60export interface GpuModelOptions {
61 device: GPUDevice;
62 ops: OpPlan;
63 air: AirGrid;
64 string: StringGrid;
65 medium: MediumFields;
66 /** Model source (.m text). */
67 source: string;
68 /** Parameter names the .m may take as arguments. */
69 paramNames: string[];
70 /** State fields the .m advances, in the order its functions return them. */
71 state: StateField[];
72 /** Grid fields one kernel may read, overriding what the device allows.
73 * Only for tests. */
74 operandBudget?: number;
75}
77/** Names the .m may take for the air grid coordinates. */
78export const AIR_GRID_NAMES = ['x', 'y', 'z'] as const;
79/** Names the .m may take for the string grid: node positions and the pin
80 * mask that terminates the ends. */
81export const STRING_GRID_NAMES = ['xs', 'pin'] as const;
82/** Names the .m may take for what the scene defines. */
83export const MEDIUM_NAMES = ['c', 'sig', 'wall', 'lineprof', 'boardprof'] as const;
85export class GpuModel {
86 readonly paramNames: string[];
87 readonly state: StateField[];
88 readonly npts: number;
89 readonly ns: number;
91 #device: GPUDevice;
92 #host: HostBuffers;
93 #initPlan: ModelPlan;
94 #stepPlan: ModelPlan;
95 /** Timestep, host-owned: it follows from the grid and the medium (a CFL
96 * condition), not from anything the user types, and it is folded into every
97 * setParams so the .m's `dt` is never left with the zero a missing
98 * parameter would default to. */
99 #dt = 0;
100 #readback: GPUBuffer;
101 /** Which function wrote the state most recently; see `read`. */
102 #lastRan: 'init' | 'step' = 'init';
104 private constructor(init: {
105 device: GPUDevice;
106 host: HostBuffers;
107 initPlan: ModelPlan;
108 stepPlan: ModelPlan;
109 readback: GPUBuffer;
110 paramNames: string[];
111 state: StateField[];
112 npts: number;
113 ns: number;
114 }) {
115 this.#device = init.device;
116 this.#host = init.host;
117 this.#initPlan = init.initPlan;
118 this.#stepPlan = init.stepPlan;
119 this.#readback = init.readback;
120 this.paramNames = init.paramNames;
121 this.state = init.state;
122 this.npts = init.npts;
123 this.ns = init.ns;
124 }
126 static async create(opts: GpuModelOptions): Promise<GpuModel> {
127 const { device, ops, air, string, medium, source, paramNames, state } = opts;
128 const npts = air.npts;
129 const ns = string.ns;
130 const sizeOf = (grid: 'air' | 'string'): number => (grid === 'air' ? npts : ns);
132 // What the .m may ask for by name. The grid geometry is exact, so a
133 // constructor reading it (`zeros(npts, 1)`) keeps a static shape.
134 const bindings: Record<string, Binding> = {
135 npts: { kind: 'const', value: npts },
136 nx: { kind: 'const', value: air.nx },
137 ny: { kind: 'const', value: air.ny },
138 nz: { kind: 'const', value: air.nz },
139 h: { kind: 'const', value: air.h },
140 ns: { kind: 'const', value: ns },
141 hs: { kind: 'const', value: string.hs },
142 Ls: { kind: 'const', value: string.Ls },
143 dt: { kind: 'param' },
144 };
145 for (const g of AIR_GRID_NAMES) bindings[g] = { kind: 'tensor', shape: [npts, 1] };
146 for (const g of STRING_GRID_NAMES) bindings[g] = { kind: 'tensor', shape: [ns, 1] };
147 for (const m of MEDIUM_NAMES) bindings[m] = { kind: 'tensor', shape: [npts, 1] };
148 for (const s of state) bindings[s.name] = { kind: 'tensor', shape: [sizeOf(s.grid), 1] };
149 for (const p of paramNames) bindings[p] = { kind: 'param' };
151 // Parsing belongs to the file, not to either function.
152 const compiled = inModel(() => new CompiledModel(source, bindings, { npts, ns }));
153 const nargout = state.length;
154 const initFn = inFunction('init', () => compiled.specialize('init', nargout));
155 const stepFn = inFunction('step', () => compiled.specialize('step', nargout));
156 compiled.finish();
158 // Both functions return the state, in order, and both feed it back into
159 // the shared buffers.
160 const feedback = state.map((s) => s.name);
162 const host = new HostBuffers(device);
163 // The host owns the state and the inputs it uploads, whether or not a
164 // given function happens to take them as arguments.
165 for (const s of state) host.ensure(s.name, sizeOf(s.grid));
166 for (const g of AIR_GRID_NAMES) host.ensure(g, npts);
167 for (const g of STRING_GRID_NAMES) host.ensure(g, ns);
168 for (const m of MEDIUM_NAMES) host.ensure(m, npts);
170 const initPlan = await inFunctionAsync('init', () =>
171 ModelPlan.create(device, ops, { fn: initFn, feedback }, host, opts.operandBudget),
172 );
173 const stepPlan = await inFunctionAsync('step', () =>
174 ModelPlan.create(device, ops, { fn: stepFn, feedback }, host, opts.operandBudget),
175 );
177 host.upload('x', air.x);
178 host.upload('y', air.y);
179 host.upload('z', air.z);
180 host.upload('xs', string.xs);
181 host.upload('pin', string.pin);
182 for (const m of MEDIUM_NAMES) host.upload(m, medium[m]);
184 const readback = device.createBuffer({
185 label: 'mgpu-readback',
186 size: 4 * Math.max(npts, ns),
187 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
188 });
190 return new GpuModel({
191 device, host, initPlan, stepPlan, readback, paramNames, state, npts, ns,
192 });
193 }
195 /** The timestep in force. Host-owned; see `#dt`. */
196 get dt(): number {
197 return this.#dt;
198 }
200 setDt(dt: number): void {
201 this.#dt = dt;
202 }
204 setParams(params: ModelParams): void {
205 const merged = { dt: this.#dt, ...params };
206 this.#initPlan.setParams(merged);
207 this.#stepPlan.setParams(merged);
208 }
210 /**
211 * Swap the medium under a running model. It is data, not code — its shape in
212 * the bindings depends only on the grid — so changing the body's geometry is
213 * five buffer writes and needs no recompile.
214 */
215 uploadMedium(medium: MediumFields): void {
216 for (const m of MEDIUM_NAMES) this.#host.upload(m, medium[m]);
217 }
219 /** Write a host-owned value directly. Lets a test set up an exact initial
220 * condition instead of going through `init`. */
221 upload(name: string, data: Float32Array): void {
222 this.#host.upload(name, data);
223 }
225 /** Run `init`, replacing the state. */
226 init(): void {
227 const enc = this.#device.createCommandEncoder({ label: 'mgpu-init' });
228 this.#initPlan.encodeSteps(enc, 1);
229 this.#device.queue.submit([enc.finish()]);
230 this.#lastRan = 'init';
231 }
233 /**
234 * Advance `steps` timesteps. Synchronous — this only records commands and
235 * submits them; nothing is read back and nothing is awaited.
236 *
237 * `after` is recorded once per step, so anything that must see every
238 * timestep (the microphone) rides along in the same submission.
239 */
240 step(steps = 1, after?: (encoder: GPUCommandEncoder) => void): void {
241 const enc = this.#device.createCommandEncoder({ label: 'mgpu-step' });
242 this.#stepPlan.encodeSteps(enc, steps, after);
243 this.#device.queue.submit([enc.finish()]);
244 this.#lastRan = 'step';
245 }
247 /**
248 * The buffer currently holding a named value. A field the .m computes is
249 * produced by both functions, into separate buffers (only the state is
250 * shared), so this resolves to whichever function ran most recently — which
251 * is what makes the first frame show the initial state rather than an
252 * unwritten buffer.
253 */
254 #locate(name: string): { buffer: GPUBuffer; count: number } | null {
255 const [first, second] =
256 this.#lastRan === 'init'
257 ? [this.#initPlan, this.#stepPlan]
258 : [this.#stepPlan, this.#initPlan];
259 const buffer = first.buffer(name) ?? second.buffer(name);
260 const count = first.elementCount(name) ?? second.elementCount(name);
261 if (!buffer || count === undefined) return null;
262 return { buffer, count };
263 }
265 /** The GPU buffer a named value would be read from right now. */
266 valueBuffer(name: string): GPUBuffer | null {
267 return this.#locate(name)?.buffer ?? null;
268 }
270 /**
271 * The buffer a host-owned field lives in — the state between calls, or an
272 * input like the wall mask.
273 *
274 * This is what the renderer binds, and it must be this rather than
275 * `valueBuffer`: a bind group is built once and holds a particular buffer,
276 * while `init` and `step` write their outputs into buffers of their own and
277 * only agree here, where their feedback copies land. Binding either
278 * function's private buffer would draw a stale field for half the run.
279 */
280 stateBuffer(name: string): GPUBuffer | null {
281 return this.#host.get(name)?.buffer ?? null;
282 }
284 /** Read a named value back to the CPU. The only await in the whole loop. */
285 async read(name: string): Promise<Float32Array> {
286 const located = this.#locate(name);
287 if (!located) throw new Error(`read: the model has no value named '${name}'`);
288 const { buffer, count } = located;
289 const enc = this.#device.createCommandEncoder({ label: `mgpu-read-${name}` });
290 enc.copyBufferToBuffer(buffer, 0, this.#readback, 0, 4 * count);
291 this.#device.queue.submit([enc.finish()]);
292 await this.#readback.mapAsync(GPUMapMode.READ, 0, 4 * count);
293 const out = new Float32Array(this.#readback.getMappedRange(0, 4 * count).slice(0));
294 this.#readback.unmap();
295 return out;
296 }
298 /** What the .m compiled to, for display. */
299 describe(): { init: string[]; step: string[] } {
300 return { init: this.#initPlan.describe(), step: this.#stepPlan.describe() };
301 }
303 destroy(): void {
304 this.#initPlan.destroy();
305 this.#stepPlan.destroy();
306 this.#host.destroy();
307 this.#readback.destroy();
308 }
309}