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: about 12 seconds of audio at a typical playback
18 * rate, and 1 MB of GPU memory. Recording stops when it is full rather than
19 * wrapping, so what you hear always starts where the run did. */
20export const TRACE_CAPACITY = 1 << 18;
22const SHADER = `
23struct Probe {
24 index: u32,
25 capacity: u32,
26};
28@group(0) @binding(0) var<storage, read_write> trace: array<f32>;
29@group(0) @binding(1) var<storage, read_write> head: array<u32>;
30@group(0) @binding(2) var<storage, read> field: array<f32>;
31@group(0) @binding(3) var<uniform> probe: Probe;
33// One invocation, so the read-modify-write of the head needs no atomic.
34@compute @workgroup_size(1)
35fn main() {
36 let i = head[0];
37 if (i < probe.capacity) {
38 trace[i] = field[probe.index];
39 head[0] = i + 1u;
40 }
41}
42`;
44export interface RecorderOptions {
45 device: GPUDevice;
46 /** The pressure buffer to sample — the host-owned one, which is what both
47 * `init` and `step` leave their result in. */
48 field: GPUBuffer;
49 nx: number;
50 ny: number;
51}
53export class Recorder {
54 readonly capacity = TRACE_CAPACITY;
56 #device: GPUDevice;
57 #pipeline: GPUComputePipeline | null = null;
58 #layout: GPUBindGroupLayout;
59 #bindGroup: GPUBindGroup | null = null;
60 #trace: GPUBuffer;
61 #head: GPUBuffer;
62 #probe: GPUBuffer;
63 #readback: GPUBuffer;
64 #field: GPUBuffer;
65 #nx: number;
66 #ny: number;
67 /** Samples written since the last clear, as far as the host knows. Counted
68 * here rather than read back from the GPU: the dispatch runs once per step
69 * and the host knows exactly how many steps it asked for. */
70 #count = 0;
71 #reading = false;
73 private constructor(init: {
74 device: GPUDevice;
75 layout: GPUBindGroupLayout;
76 trace: GPUBuffer;
77 head: GPUBuffer;
78 probe: GPUBuffer;
79 readback: GPUBuffer;
80 field: GPUBuffer;
81 nx: number;
82 ny: number;
83 }) {
84 this.#device = init.device;
85 this.#layout = init.layout;
86 this.#trace = init.trace;
87 this.#head = init.head;
88 this.#probe = init.probe;
89 this.#readback = init.readback;
90 this.#field = init.field;
91 this.#nx = init.nx;
92 this.#ny = init.ny;
93 }
95 static async create(opts: RecorderOptions): Promise<Recorder> {
96 const { device } = opts;
97 const layout = device.createBindGroupLayout({
98 label: 'recorder',
99 entries: [
100 { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
101 { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } },
102 {
103 binding: 2,
104 visibility: GPUShaderStage.COMPUTE,
105 buffer: { type: 'read-only-storage' },
106 },
107 { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform' } },
108 ],
109 });
110 const trace = device.createBuffer({
111 label: 'recorder-trace',
112 size: 4 * TRACE_CAPACITY,
113 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
114 });
115 const head = device.createBuffer({
116 label: 'recorder-head',
117 size: 4,
118 usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
119 });
120 const probe = device.createBuffer({
121 label: 'recorder-probe',
122 size: 8,
123 usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
124 });
125 const readback = device.createBuffer({
126 label: 'recorder-readback',
127 size: 4 * TRACE_CAPACITY,
128 usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
129 });
131 const rec = new Recorder({
132 device, layout, trace, head, probe, readback,
133 field: opts.field, nx: opts.nx, ny: opts.ny,
134 });
135 rec.#pipeline = await device.createComputePipelineAsync({
136 label: 'recorder',
137 layout: device.createPipelineLayout({ bindGroupLayouts: [layout] }),
138 compute: {
139 module: device.createShaderModule({ code: SHADER, label: 'recorder' }),
140 entryPoint: 'main',
141 },
142 });
143 rec.#bindGroup = device.createBindGroup({
144 layout,
145 entries: [
146 { binding: 0, resource: { buffer: trace } },
147 { binding: 1, resource: { buffer: head } },
148 { binding: 2, resource: { buffer: opts.field } },
149 { binding: 3, resource: { buffer: probe } },
150 ],
151 });
152 rec.clear();
153 return rec;
154 }
156 /** Samples recorded so far. Stops rising once the trace is full. */
157 get count(): number {
158 return Math.min(this.#count, this.capacity);
159 }
161 get full(): boolean {
162 return this.#count >= this.capacity;
163 }
165 /** Put the microphone at the grid point nearest (x, y). Free: the probe
166 * index is a uniform, so moving it disturbs neither the run nor the
167 * recording already made. */
168 setProbe(ix: number, iy: number): void {
169 const cx = Math.max(0, Math.min(this.#nx - 1, Math.round(ix)));
170 const cy = Math.max(0, Math.min(this.#ny - 1, Math.round(iy)));
171 this.#device.queue.writeBuffer(
172 this.#probe,
173 0,
174 new Uint32Array([cx + this.#nx * cy, this.capacity]),
175 );
176 }
178 /** Start again from an empty trace. */
179 clear(): void {
180 this.#count = 0;
181 this.#device.queue.writeBuffer(this.#head, 0, new Uint32Array([0]));
182 }
184 /** Record one sample. Called once per timestep, inside the step's own
185 * submission, so no extra work crosses to the host. */
186 encode(encoder: GPUCommandEncoder): void {
187 if (!this.#pipeline || !this.#bindGroup || this.full) return;
188 const pass = encoder.beginComputePass({ label: 'recorder' });
189 pass.setPipeline(this.#pipeline);
190 pass.setBindGroup(0, this.#bindGroup);
191 pass.dispatchWorkgroups(1);
192 pass.end();
193 this.#count++;
194 }
196 /** The recorded trace. The only readback the microphone ever does. */
197 async read(): Promise<Float32Array> {
198 const n = this.count;
199 if (n === 0) return new Float32Array(0);
200 if (this.#reading) throw new Error('a trace readback is already in flight');
201 this.#reading = true;
202 try {
203 const enc = this.#device.createCommandEncoder({ label: 'recorder-read' });
204 enc.copyBufferToBuffer(this.#trace, 0, this.#readback, 0, 4 * n);
205 this.#device.queue.submit([enc.finish()]);
206 await this.#readback.mapAsync(GPUMapMode.READ, 0, 4 * n);
207 const out = new Float32Array(this.#readback.getMappedRange(0, 4 * n).slice(0));
208 this.#readback.unmap();
209 return out;
210 } finally {
211 this.#reading = false;
212 }
213 }
215 /** The pressure buffer this microphone listens to. */
216 get field(): GPUBuffer {
217 return this.#field;
218 }
220 destroy(): void {
221 this.#trace.destroy();
222 this.#head.destroy();
223 this.#probe.destroy();
224 this.#readback.destroy();
225 }
226}