1/**
2 * A microphone: the pressure at one grid point, sampled every timestep.
3 *
4 * The obvious implementation — read the field back and pick out one number —
5 * costs a GPU-to-CPU round trip per step, which is more than the step itself.
6 * So the trace is written on the GPU instead, by a one-thread dispatch that
7 * runs after each step and appends `p` at the probe point to a buffer. The
8 * whole trace comes back to the CPU once, when there is something to listen
9 * to.
10 *
11 * One sample per timestep is the natural rate: it is every value the
12 * simulation has, and nothing is being resampled or interpolated on the way
13 * in. What that means in seconds is decided at playback (src/audio/play.ts),
14 * because the simulation has no seconds in it — only model time.
15 */
17/** Samples the trace holds by default: about 9 seconds of audio at the
18 * browser's default timestep, and 4 MB of GPU memory. Recording stops when
19 * it is full rather than wrapping, so what you hear always starts where the
20 * excitation did. The headless renderer passes its own capacity, sized from
21 * the requested duration. */
22export const TRACE_CAPACITY = 1 << 20;
24const SHADER = `
25struct Probe {
26 index: u32,
27 capacity: u32,
28};
30@group(0) @binding(0) var<storage, read_write> trace: array<f32>;
31@group(0) @binding(1) var<storage, read_write> head: array<u32>;
32@group(0) @binding(2) var<storage, read> field: array<f32>;
33@group(0) @binding(3) var<uniform> probe: Probe;
35// One invocation, so the read-modify-write of the head needs no atomic.
36@compute @workgroup_size(1)
37fn main() {
38 let i = head[0];
39 if (i < probe.capacity) {
40 trace[i] = field[probe.index];
41 head[0] = i + 1u;
42 }
43}
44`;
46export interface RecorderOptions {
47 device: GPUDevice;
48 /** The pressure buffer to sample — the host-owned one, which is what both
49 * `init` and `step` leave their result in. */
50 field: GPUBuffer;
51 nx: number;
52 ny: number;
53 nz: number;
54 /** Samples the trace holds; defaults to TRACE_CAPACITY. */
55 capacity?: number;
56}
58export class Recorder {
59 readonly capacity: number;
61 #device: GPUDevice;
62 #pipeline: GPUComputePipeline | null = null;
63 #layout: GPUBindGroupLayout;
64 #bindGroup: GPUBindGroup | null = null;
65 #trace: GPUBuffer;
66 #head: GPUBuffer;
67 #probe: GPUBuffer;
68 #readback: GPUBuffer;
69 #field: GPUBuffer;
70 #nx: number;
71 #ny: number;
72 #nz: number;
73 /** Samples written since the last clear, as far as the host knows. Counted
74 * here rather than read back from the GPU: the dispatch runs once per step
75 * and the host knows exactly how many steps it asked for. */
76 #count = 0;
77 #reading = false;
79 private constructor(init: {
80 device: GPUDevice;
81 layout: GPUBindGroupLayout;
82 trace: GPUBuffer;
83 head: GPUBuffer;
84 probe: GPUBuffer;
85 readback: GPUBuffer;
86 field: GPUBuffer;
87 nx: number;
88 ny: number;
89 nz: number;
90 capacity: number;
91 }) {
92 this.capacity = init.capacity;
93 this.#device = init.device;
94 this.#layout = init.layout;
95 this.#trace = init.trace;
96 this.#head = init.head;
97 this.#probe = init.probe;
98 this.#readback = init.readback;
99 this.#field = init.field;
100 this.#nx = init.nx;
101 this.#ny = init.ny;
102 this.#nz = init.nz;
103 }
105 static async create(opts: RecorderOptions): Promise<Recorder> {
106 const { device } = opts;
107 const capacity = opts.capacity ?? TRACE_CAPACITY;
108 const layout = device.createBindGroupLayout({
109 label: 'recorder',
110 entries: [
111 { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
112 { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
113 {
114 binding: 2,
115 visibility: GPUShaderStage.COMPUTE,
116 buffer: { type: 'read-only-storage' },
117 },
118 { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform' } },
119 ],
120 });
121 const trace = device.createBuffer({
122 label: 'recorder-trace',
123 size: 4 * capacity,
124 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
125 });
126 const head = device.createBuffer({
127 label: 'recorder-head',
128 size: 4,
129 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
130 });
131 const probe = device.createBuffer({
132 label: 'recorder-probe',
133 size: 8,
134 usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
135 });
136 const readback = device.createBuffer({
137 label: 'recorder-readback',
138 size: 4 * capacity,
139 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
140 });
142 const rec = new Recorder({
143 device, layout, trace, head, probe, readback,
144 field: opts.field, nx: opts.nx, ny: opts.ny, nz: opts.nz, capacity,
145 });
146 rec.#pipeline = await device.createComputePipelineAsync({
147 label: 'recorder',
148 layout: device.createPipelineLayout({ bindGroupLayouts: [layout] }),
149 compute: {
150 module: device.createShaderModule({ code: SHADER, label: 'recorder' }),
151 entryPoint: 'main',
152 },
153 });
154 rec.#bindGroup = device.createBindGroup({
155 layout,
156 entries: [
157 { binding: 0, resource: { buffer: trace } },
158 { binding: 1, resource: { buffer: head } },
159 { binding: 2, resource: { buffer: opts.field } },
160 { binding: 3, resource: { buffer: probe } },
161 ],
162 });
163 rec.clear();
164 return rec;
165 }
167 /** Samples recorded so far. Stops rising once the trace is full. */
168 get count(): number {
169 return Math.min(this.#count, this.capacity);
170 }
172 get full(): boolean {
173 return this.#count >= this.capacity;
174 }
176 /** Put the microphone at the grid point nearest (ix, iy, iz). Free: the
177 * probe index is a uniform, so moving it disturbs neither the run nor the
178 * recording already made. */
179 setProbe(ix: number, iy: number, iz: number): void {
180 const cx = Math.max(0, Math.min(this.#nx - 1, Math.round(ix)));
181 const cy = Math.max(0, Math.min(this.#ny - 1, Math.round(iy)));
182 const cz = Math.max(0, Math.min(this.#nz - 1, Math.round(iz)));
183 this.#device.queue.writeBuffer(
184 this.#probe,
185 0,
186 new Uint32Array([cx + this.#nx * (cy + this.#ny * cz), this.capacity]),
187 );
188 }
190 /** Start again from an empty trace. */
191 clear(): void {
192 this.#count = 0;
193 this.#device.queue.writeBuffer(this.#head, 0, new Uint32Array([0]));
194 }
196 /** Record one sample. Called once per timestep, inside the step's own
197 * submission, so no extra work crosses to the host. */
198 encode(encoder: GPUCommandEncoder): void {
199 if (!this.#pipeline || !this.#bindGroup || this.full) return;
200 const pass = encoder.beginComputePass({ label: 'recorder' });
201 pass.setPipeline(this.#pipeline);
202 pass.setBindGroup(0, this.#bindGroup);
203 pass.dispatchWorkgroups(1);
204 pass.end();
205 this.#count++;
206 }
208 /** The recorded trace. The only readback the microphone ever does. */
209 async read(): Promise<Float32Array> {
210 const n = this.count;
211 if (n === 0) return new Float32Array(0);
212 if (this.#reading) throw new Error('a trace readback is already in flight');
213 this.#reading = true;
214 try {
215 const enc = this.#device.createCommandEncoder({ label: 'recorder-read' });
216 enc.copyBufferToBuffer(this.#trace, 0, this.#readback, 0, 4 * n);
217 this.#device.queue.submit([enc.finish()]);
218 await this.#readback.mapAsync(GPUMapMode.READ, 0, 4 * n);
219 const out = new Float32Array(this.#readback.getMappedRange(0, 4 * n).slice(0));
220 this.#readback.unmap();
221 return out;
222 } finally {
223 this.#reading = false;
224 }
225 }
227 /** The pressure buffer this microphone listens to. */
228 get field(): GPUBuffer {
229 return this.#field;
230 }
232 destroy(): void {
233 this.#trace.destroy();
234 this.#head.destroy();
235 this.#probe.destroy();
236 this.#readback.destroy();
237 }
238}