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