1/**
2 * The solver: a second-order leapfrog for the acoustic wave equation on a
3 * cubic grid,
4 *
5 * p_tt + 2*sig*p_t = c(x)^2 * lap(p) + s(x, t)
6 *
7 * Pressure only, at constant density, so the medium is two fields: the sound
8 * speed c and the absorption sig, both supplied by the scene. Centring the
9 * second time derivative and the damping on step n gives an explicit update,
10 * which is one compute dispatch per timestep.
11 *
12 * Two things are worth explaining.
13 *
14 * **The update is in place.** A step reads p at its six neighbours but reads
15 * the previous field pm only at its own index, so a thread may overwrite
16 * pm[i] with the new value: no other thread will read it. That leaves two
17 * pressure buffers instead of three, which matters at 192^3 where each one is
18 * 28 MB, and it removes the buffer-to-buffer copy per step. The roles swap
19 * every step, so which buffer holds the current field depends on the parity of
20 * the step count; `pressure` reports it and the renderer follows.
21 *
22 * **Time is a uniform read at a dynamic offset.** A frame's worth of steps is
23 * recorded into one command encoder, so nothing written between submits can
24 * change inside it, and a clock uploaded per frame would stand still for the
25 * whole batch. Rather than carrying time as a grid field (the flat sibling's
26 * answer, forced there by its compiler), each step reads its own 64-byte slice
27 * of one parameter buffer through a dynamic offset. The whole batch's
28 * parameters are written in one call before the pass.
29 */
30import type { Grid } from './grid.ts';
32/** Bytes of parameters per step. Padded to the uniform dynamic-offset
33 * alignment, which WebGPU guarantees to be at most 256. */
34const PARAM_STRIDE = 256;
36/** Steps that fit in the parameter buffer, and so in one submit. */
37export const MAX_STEPS_PER_FRAME = 64;
39/** Workgroups are reduced into this many partial maxima before readback. */
40const REDUCE_GROUPS = 64;
42/** Samples the microphone trace can hold: 4 MB of f32, a few seconds of
43 * audio at the rates the default timestep implies. */
44export const MIC_CAPACITY = 1 << 20;
46const STEP_SHADER = `
47struct Params {
48 n: u32,
49 h: f32,
50 dt: f32,
51 t: f32,
52 L: f32,
53 f: f32,
54 t0: f32,
55 tw: f32,
56 cw: f32,
57 x0: f32,
58 y0: f32,
59 z0: f32,
60 w: f32,
61 point: f32,
62};
64@group(0) @binding(0) var<uniform> P: Params;
65@group(1) @binding(0) var<storage, read> p: array<f32>;
66@group(1) @binding(1) var<storage, read_write> pm: array<f32>;
67@group(1) @binding(2) var<storage, read> cs: array<f32>;
68@group(1) @binding(3) var<storage, read> sg: array<f32>;
70// Outside the domain the field is zero. The absorbing layer is meant to have
71// swallowed the wave long before it reaches here.
72fn at(ix: i32, iy: i32, iz: i32) -> f32 {
73 let n = i32(P.n);
74 if (ix < 0 || iy < 0 || iz < 0 || ix >= n || iy >= n || iz >= n) { return 0.0; }
75 return p[u32(ix + n * iy + n * n * iz)];
76}
78@compute @workgroup_size(8, 8, 4)
79fn main(@builtin(global_invocation_id) gid: vec3u) {
80 let n = P.n;
81 if (gid.x >= n || gid.y >= n || gid.z >= n) { return; }
82 let i = gid.x + n * gid.y + n * n * gid.z;
83 let ix = i32(gid.x);
84 let iy = i32(gid.y);
85 let iz = i32(gid.z);
87 let pc = p[i];
88 let lap = (at(ix + 1, iy, iz) + at(ix - 1, iy, iz)
89 + at(ix, iy + 1, iz) + at(ix, iy - 1, iz)
90 + at(ix, iy, iz + 1) + at(ix, iy, iz - 1)
91 - 6.0 * pc) / (P.h * P.h);
93 let x = -0.5 * P.L + (f32(ix) + 0.5) * P.h;
94 let y = -0.5 * P.L + (f32(iy) + 0.5) * P.h;
95 let z = -0.5 * P.L + (f32(iz) + 0.5) * P.h;
97 // The source. \`cw\` blends between a Gaussian pulse (0) and a wave that
98 // turns on smoothly and stays on (1); \`point\` blends between a planar
99 // source spanning the grid in y and z, whose far field is a plane wave, and
100 // a point source at (x0, y0, z0).
101 //
102 // The om^2 is a choice of units, not a physical amplitude: a body force of
103 // fixed strength drives a response falling off as 1/om^2, so without it the
104 // field would shrink every time the frequency slider went up.
105 let u = (P.t - P.t0) / P.tw;
106 let env = (1.0 - P.cw) * exp(-u * u) + P.cw * 0.5 * (1.0 + tanh(u));
107 let gx = (x - P.x0) / P.w;
108 let gy = P.point * (y - P.y0) / P.w;
109 let gz = P.point * (z - P.z0) / P.w;
110 let om = 6.283185307179586 * P.f;
111 let s = om * om * env * sin(om * (P.t - P.t0)) * exp(-(gx * gx + gy * gy + gz * gz));
113 // One step. The damping is what the absorbing layer acts through: sig is
114 // zero over the interior, so there this is the plain leapfrog update.
115 let sd = sg[i] * P.dt;
116 let cd = cs[i] * P.dt;
117 pm[i] = (2.0 * pc - (1.0 - sd) * pm[i] + cd * cd * lap + P.dt * P.dt * s) / (1.0 + sd);
118}
119`;
121// The microphone: the pressure at one grid point, appended to a trace by a
122// one-thread dispatch that runs after each step inside the same command
123// stream. The obvious implementation — reading the field back and picking out
124// one number — costs a GPU-to-CPU round trip per step and would be slower
125// than the step; this way the whole trace comes back once, when there is
126// something to play. One thread and in-order dispatches mean the plain
127// read-modify-write on the counter is safe.
128const MIC_SHADER = `
129struct MicInfo { probe: u32, cap: u32 };
130@group(0) @binding(0) var<uniform> M: MicInfo;
131@group(0) @binding(1) var<storage, read> p: array<f32>;
132@group(0) @binding(2) var<storage, read_write> trace: array<f32>;
133@group(0) @binding(3) var<storage, read_write> count: array<u32>;
135@compute @workgroup_size(1)
136fn main() {
137 let k = count[0];
138 if (k < M.cap) {
139 trace[k] = p[M.probe];
140 count[0] = k + 1u;
141 }
142}
143`;
145const REDUCE_SHADER = `
146@group(0) @binding(0) var<uniform> npts: u32;
147@group(0) @binding(1) var<storage, read> p: array<f32>;
148@group(0) @binding(2) var<storage, read_write> out: array<f32>;
150var<workgroup> sh: array<f32, 256>;
152// Both loops are given uniform trip counts on purpose: a workgroupBarrier may
153// only be reached in uniform control flow, and a loop whose exit depends on
154// the thread index taints everything after it.
155@compute @workgroup_size(256)
156fn main(@builtin(global_invocation_id) gid: vec3u,
157 @builtin(local_invocation_id) lid: vec3u,
158 @builtin(workgroup_id) wid: vec3u) {
159 let stride = 256u * ${REDUCE_GROUPS}u;
160 let per = (npts + stride - 1u) / stride;
161 var m = 0.0;
162 for (var k = 0u; k < per; k = k + 1u) {
163 let i = gid.x + k * stride;
164 if (i < npts) { m = max(m, abs(p[i])); }
165 }
166 sh[lid.x] = m;
167 workgroupBarrier();
168 for (var s = 128u; s > 0u; s = s >> 1u) {
169 if (lid.x < s) { sh[lid.x] = max(sh[lid.x], sh[lid.x + s]); }
170 workgroupBarrier();
171 }
172 if (lid.x == 0u) { out[wid.x] = sh[0]; }
173}
174`;
176/** Everything the source term needs, in SI. */
177export interface SourceParams {
178 f: number;
179 cycles: number;
180 cw: number;
181 point: number;
182 x0: number;
183 y0: number;
184 z0: number;
185 w: number;
186}
188export class Sim {
189 readonly grid: Grid;
190 /** Model time, seconds. */
191 t = 0;
192 steps = 0;
193 dt: number;
194 source: SourceParams;
196 #device: GPUDevice;
197 #pa: GPUBuffer;
198 #pb: GPUBuffer;
199 #c: GPUBuffer;
200 #sig: GPUBuffer;
201 #params: GPUBuffer;
202 #paramsHost = new ArrayBuffer(MAX_STEPS_PER_FRAME * PARAM_STRIDE);
203 #paramsBG: GPUBindGroup;
204 #fieldBG: [GPUBindGroup, GPUBindGroup];
205 #pipeline: GPUComputePipeline;
206 /** Index of the field bind group to use for the next step. It is also which
207 * buffer currently holds the field: 0 means `pa`. */
208 #next = 0;
210 #reducePipeline: GPUComputePipeline;
211 #reduceBG: [GPUBindGroup, GPUBindGroup];
212 #partials: GPUBuffer;
213 #readback: GPUBuffer;
214 #reading = false;
216 /** Samples recorded so far. A host-side mirror of the GPU counter: both
217 * add one per step until the capacity, so they agree exactly. */
218 recorded = 0;
219 #micPipeline: GPUComputePipeline;
220 #micBG: [GPUBindGroup, GPUBindGroup];
221 #micUniform: GPUBuffer;
222 #trace: GPUBuffer;
223 #micCount: GPUBuffer;
224 #micRB: GPUBuffer;
225 #micReading = false;
227 constructor(
228 device: GPUDevice,
229 grid: Grid,
230 medium: { c: Float32Array<ArrayBuffer>; sig: Float32Array<ArrayBuffer> },
231 source: SourceParams,
232 dt: number,
233 ) {
234 this.#device = device;
235 this.grid = grid;
236 this.source = source;
237 this.dt = dt;
239 const bytes = grid.npts * 4;
240 const field = (label: string) =>
241 device.createBuffer({
242 label,
243 size: bytes,
244 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
245 });
246 this.#pa = field('p-a');
247 this.#pb = field('p-b');
248 this.#c = field('speed');
249 this.#sig = field('absorption');
250 device.queue.writeBuffer(this.#c, 0, medium.c);
251 device.queue.writeBuffer(this.#sig, 0, medium.sig);
253 this.#params = device.createBuffer({
254 label: 'step-params',
255 size: this.#paramsHost.byteLength,
256 usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
257 });
259 const module = device.createShaderModule({ code: STEP_SHADER, label: 'leapfrog3d' });
260 const paramsLayout = device.createBindGroupLayout({
261 entries: [
262 {
263 binding: 0,
264 visibility: GPUShaderStage.COMPUTE,
265 buffer: { type: 'uniform', hasDynamicOffset: true, minBindingSize: 64 },
266 },
267 ],
268 });
269 const fieldLayout = device.createBindGroupLayout({
270 entries: [
271 { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
272 { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
273 { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
274 { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
275 ],
276 });
277 this.#pipeline = device.createComputePipeline({
278 label: 'leapfrog3d',
279 layout: device.createPipelineLayout({ bindGroupLayouts: [paramsLayout, fieldLayout] }),
280 compute: { module, entryPoint: 'main' },
281 });
282 this.#paramsBG = device.createBindGroup({
283 layout: paramsLayout,
284 entries: [{ binding: 0, resource: { buffer: this.#params, size: 64 } }],
285 });
286 const pair = (read: GPUBuffer, write: GPUBuffer) =>
287 device.createBindGroup({
288 layout: fieldLayout,
289 entries: [
290 { binding: 0, resource: { buffer: read } },
291 { binding: 1, resource: { buffer: write } },
292 { binding: 2, resource: { buffer: this.#c } },
293 { binding: 3, resource: { buffer: this.#sig } },
294 ],
295 });
296 this.#fieldBG = [pair(this.#pa, this.#pb), pair(this.#pb, this.#pa)];
298 // The colour scale wants max |p| over the whole field, which at 128^3 is
299 // 8 MB to read back. A workgroup reduction turns it into 64 floats first.
300 const nptsBuf = device.createBuffer({
301 size: 4,
302 usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
303 });
304 device.queue.writeBuffer(nptsBuf, 0, new Uint32Array([grid.npts]));
305 this.#partials = device.createBuffer({
306 size: REDUCE_GROUPS * 4,
307 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
308 });
309 this.#readback = device.createBuffer({
310 size: REDUCE_GROUPS * 4,
311 usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
312 });
313 const reduceLayout = device.createBindGroupLayout({
314 entries: [
315 { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform' } },
316 { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
317 { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
318 ],
319 });
320 this.#reducePipeline = device.createComputePipeline({
321 label: 'maxabs',
322 layout: device.createPipelineLayout({ bindGroupLayouts: [reduceLayout] }),
323 compute: {
324 module: device.createShaderModule({ code: REDUCE_SHADER, label: 'maxabs' }),
325 entryPoint: 'main',
326 },
327 });
328 const red = (p: GPUBuffer) =>
329 device.createBindGroup({
330 layout: reduceLayout,
331 entries: [
332 { binding: 0, resource: { buffer: nptsBuf } },
333 { binding: 1, resource: { buffer: p } },
334 { binding: 2, resource: { buffer: this.#partials } },
335 ],
336 });
337 this.#reduceBG = [red(this.#pa), red(this.#pb)];
339 // The microphone. All buffers start zeroed, so the counter needs no init.
340 this.#micUniform = device.createBuffer({
341 label: 'mic-info',
342 size: 8,
343 usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
344 });
345 this.#trace = device.createBuffer({
346 label: 'mic-trace',
347 size: MIC_CAPACITY * 4,
348 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
349 });
350 this.#micCount = device.createBuffer({
351 label: 'mic-count',
352 size: 4,
353 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
354 });
355 this.#micRB = device.createBuffer({
356 label: 'mic-readback',
357 size: MIC_CAPACITY * 4,
358 usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
359 });
360 const micLayout = device.createBindGroupLayout({
361 entries: [
362 { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform' } },
363 { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
364 { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
365 { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
366 ],
367 });
368 this.#micPipeline = device.createComputePipeline({
369 label: 'microphone',
370 layout: device.createPipelineLayout({ bindGroupLayouts: [micLayout] }),
371 compute: {
372 module: device.createShaderModule({ code: MIC_SHADER, label: 'microphone' }),
373 entryPoint: 'main',
374 },
375 });
376 const micOn = (p: GPUBuffer) =>
377 device.createBindGroup({
378 layout: micLayout,
379 entries: [
380 { binding: 0, resource: { buffer: this.#micUniform } },
381 { binding: 1, resource: { buffer: p } },
382 { binding: 2, resource: { buffer: this.#trace } },
383 { binding: 3, resource: { buffer: this.#micCount } },
384 ],
385 });
386 this.#micBG = [micOn(this.#pa), micOn(this.#pb)];
387 this.setProbe(0, 0, 0);
388 }
390 /** Both pressure buffers; `pressureIndex` says which holds the field now. */
391 get pressures(): [GPUBuffer, GPUBuffer] {
392 return [this.#pa, this.#pb];
393 }
394 get pressureIndex(): number {
395 return this.#next;
396 }
397 get speed(): GPUBuffer {
398 return this.#c;
399 }
401 /** Point the microphone at a grid cell. Does not restart the trace: moving
402 * the microphone during a run is a microphone that moved. */
403 setProbe(ix: number, iy: number, iz: number): void {
404 const n = this.grid.n;
405 const cl = (i: number) => Math.min(n - 1, Math.max(0, i));
406 const probe = cl(ix) + n * cl(iy) + n * n * cl(iz);
407 this.#device.queue.writeBuffer(this.#micUniform, 0, new Uint32Array([probe, MIC_CAPACITY]));
408 }
410 /** Start the recording over. Called on restart, and whenever dt changes:
411 * the trace is one sample per step, and two timesteps would be two sample
412 * rates in one buffer. */
413 resetTrace(): void {
414 this.#device.queue.writeBuffer(this.#micCount, 0, new Uint32Array([0]));
415 this.recorded = 0;
416 }
418 /** The recorded trace, back from the GPU. Null if a read is in flight. */
419 async readTrace(): Promise<Float32Array | null> {
420 if (this.#micReading) return null;
421 if (this.recorded === 0) return new Float32Array(0);
422 this.#micReading = true;
423 try {
424 const bytes = this.recorded * 4;
425 const enc = this.#device.createCommandEncoder({ label: 'mic-read' });
426 enc.copyBufferToBuffer(this.#trace, 0, this.#micRB, 0, bytes);
427 this.#device.queue.submit([enc.finish()]);
428 await this.#micRB.mapAsync(GPUMapMode.READ, 0, bytes);
429 const out = new Float32Array(this.#micRB.getMappedRange(0, bytes).slice(0));
430 this.#micRB.unmap();
431 return out;
432 } finally {
433 this.#micReading = false;
434 }
435 }
437 /** Back to a silent grid at t = 0. */
438 restart(): void {
439 const enc = this.#device.createCommandEncoder();
440 enc.clearBuffer(this.#pa);
441 enc.clearBuffer(this.#pb);
442 this.#device.queue.submit([enc.finish()]);
443 this.t = 0;
444 this.steps = 0;
445 this.#next = 0;
446 this.resetTrace();
447 }
449 /** Take `n` timesteps, all in one submit. */
450 run(n: number): void {
451 const count = Math.min(n, MAX_STEPS_PER_FRAME);
452 const g = this.grid;
453 const s = this.source;
454 // A pulse `cycles` long, delayed enough that it starts near zero.
455 const tw = s.cycles / Math.max(s.f, 1e-6);
456 const t0 = 2.5 * tw;
457 for (let k = 0; k < count; k++) {
458 const off = k * PARAM_STRIDE;
459 const u32 = new Uint32Array(this.#paramsHost, off, 1);
460 const f32 = new Float32Array(this.#paramsHost, off, 16);
461 u32[0] = g.n;
462 f32[1] = g.h;
463 f32[2] = this.dt;
464 f32[3] = this.t + k * this.dt;
465 f32[4] = g.L;
466 f32[5] = s.f;
467 f32[6] = t0;
468 f32[7] = tw;
469 f32[8] = s.cw;
470 f32[9] = s.x0;
471 f32[10] = s.y0;
472 f32[11] = s.z0;
473 f32[12] = s.w;
474 f32[13] = s.point;
475 }
476 this.#device.queue.writeBuffer(this.#params, 0, this.#paramsHost, 0, count * PARAM_STRIDE);
478 const wg = [g.n / 8, g.n / 8, g.n / 4] as const;
479 const enc = this.#device.createCommandEncoder({ label: 'steps' });
480 const pass = enc.beginComputePass();
481 for (let k = 0; k < count; k++) {
482 pass.setPipeline(this.#pipeline);
483 pass.setBindGroup(0, this.#paramsBG, [k * PARAM_STRIDE]);
484 pass.setBindGroup(1, this.#fieldBG[this.#next]);
485 pass.dispatchWorkgroups(wg[0], wg[1], wg[2]);
486 this.#next ^= 1;
487 // Record the field this step just wrote, which the toggle now points at.
488 pass.setPipeline(this.#micPipeline);
489 pass.setBindGroup(0, this.#micBG[this.#next]);
490 pass.dispatchWorkgroups(1);
491 }
492 pass.end();
493 this.#device.queue.submit([enc.finish()]);
494 this.t += count * this.dt;
495 this.steps += count;
496 this.recorded = Math.min(this.recorded + count, MIC_CAPACITY);
497 }
499 /**
500 * Largest |p| anywhere, for the colour scale. Asynchronous and
501 * self-throttling: while one request is in flight further ones return null
502 * rather than queueing up.
503 */
504 async peak(): Promise<number | null> {
505 if (this.#reading) return null;
506 this.#reading = true;
507 try {
508 const enc = this.#device.createCommandEncoder({ label: 'peak' });
509 const pass = enc.beginComputePass();
510 pass.setPipeline(this.#reducePipeline);
511 pass.setBindGroup(0, this.#reduceBG[this.pressureIndex]);
512 pass.dispatchWorkgroups(REDUCE_GROUPS);
513 pass.end();
514 enc.copyBufferToBuffer(this.#partials, 0, this.#readback, 0, REDUCE_GROUPS * 4);
515 this.#device.queue.submit([enc.finish()]);
516 await this.#readback.mapAsync(GPUMapMode.READ);
517 const v = new Float32Array(this.#readback.getMappedRange().slice(0));
518 this.#readback.unmap();
519 let m = 0;
520 for (const x of v) m = Math.max(m, x);
521 return m;
522 } finally {
523 this.#reading = false;
524 }
525 }
527 destroy(): void {
528 const own = [this.#pa, this.#pb, this.#c, this.#sig, this.#params, this.#partials];
529 for (const b of [...own, this.#micUniform, this.#trace, this.#micCount]) b.destroy();
530 // Mapping is asynchronous; destroying a buffer with a pending map is an
531 // error, so leave the readback buffers to be collected.
532 }
533}