/ proof-of-concept / cymmodes
Sign in
proof-of-concept / cymmodes
cymmodes / src / gpu / simulator.ts
299 lines · 19.2 KBCodeBlameHistory
8354482cymmodes: linear modes of a cymsim cymbal driven in the full nonlinear model, with the energy of every mode over timeJeremy Magland 1// WebGPU time loop: packs the set-up into storage buffers, compiles the kernels and runs the
2// sub-steps of a chunk of audio samples as one command buffer. Works in a browser and under Node
3// with the `webgpu` (Dawn) package, which the tests use. Relative to cymsimweb this copy adds a
4// prescribed initial state, the distributed harmonic drive, and the modal projection pass.
5import { HOLD_HEIGHT } from "../model/exciter.ts";
6import { strainEnergies } from "../model/forces.ts";
7import { DECIMATION_SPAN } from "../model/radiation.ts";
8import type { Setup } from "../model/setup.ts";
9import type { BackendOptions, ChunkResult, Diagnostics, SimulatorBackend } from "../model/simulator.ts";
10import { CHUNK_UNIFORM_BYTES, ELEM_STRIDE, HIT_HEADER, HITU_HEADER, NODE_STRIDE, SYNC_HEADER, WORKGROUP_NODES, shaderSource } from "./shaders.ts";
12// evaluated lazily: under Node the WebGPU globals are installed after this module is imported
13const storageUsage = () => GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC;
15export class GpuSimulator implements SimulatorBackend {
16 readonly kind = "gpu" as const;
17 readonly setup: Setup; readonly device: GPUDevice; readonly opts: BackendOptions;
18 samplesDone = 0;
19 private step = 0;
20 private readonly G: number; private readonly HF: number; private readonly HU: number;
21 private readonly buffers: Record<string, GPUBuffer> = {};
22 private readonly pipelines: Record<string, GPUComputePipeline> = {};
23 private readonly bindGroups: Record<string, GPUBindGroup> = {};
24 private readonly chunkCapacity: number; readonly movieEvery: number; private readonly framesCapacity: number;
25 private readonly projEvery: number; private readonly nModes: number; private readonly projCapacity: number;
26 private readonly tracing: boolean;
27 private readonly stagingPressure: GPUBuffer; private readonly stagingState: GPUBuffer; private readonly stagingHits: GPUBuffer; private readonly stagingHitsU: GPUBuffer;
28 private readonly stagingMovie: GPUBuffer | null; private readonly stagingProj: GPUBuffer | null; private readonly stagingTrace: GPUBuffer | null;
30 constructor(device: GPUDevice, setup: Setup, chunkCapacity = 8192, opts: BackendOptions = {}) {
31 this.device = device; this.setup = setup; this.opts = opts;
32 const movieEvery = opts.movieEvery ?? 0;
33 this.movieEvery = movieEvery;
34 const Nn = setup.mesh.n;
35 if (movieEvery > 0) chunkCapacity = Math.max(128, Math.min(chunkCapacity, Math.floor((64e6 / (4 * Nn)) * movieEvery)));
36 this.chunkCapacity = chunkCapacity;
37 this.framesCapacity = movieEvery > 0 ? Math.ceil(chunkCapacity / movieEvery) + 2 : 0;
38 this.projEvery = opts.projection ? opts.projection.every : 0;
39 this.nModes = opts.projection ? opts.projection.nModes : 0;
40 this.projCapacity = this.projEvery > 0 ? Math.ceil(chunkCapacity / this.projEvery) + 2 : 0;
41 this.tracing = !!opts.projection && opts.traceMode !== undefined;
42 const { mesh, shell: sh, schedule: sc, radiation: rad, substeps: M, dt } = setup;
43 const N = mesh.n, E = mesh.e, K = sc.K, P = sc.P, W = DECIMATION_SPAN, D = rad.ringLength;
44 this.G = Math.ceil(N / WORKGROUP_NODES);
45 this.HF = HIT_HEADER + 5 * P; this.HU = HITU_HEADER + 3 * P;
46 const G = this.G, HF = this.HF, HU = this.HU;
48 // ---- pack the set-up
49 const elemF = new Float32Array(ELEM_STRIDE * E + W * M);
50 for (let e = 0; e < E; e++) {
51 const b = ELEM_STRIDE * e;
52 for (let i = 0; i < 3; i++) { elemF[b + i] = mesh.Gx[3 * e + i]; elemF[b + 3 + i] = mesh.Gy[3 * e + i]; }
53 elemF[b + 6] = mesh.area[e]; elemF[b + 7] = sh.z0x[e]; elemF[b + 8] = sh.z0y[e];
54 elemF[b + 9] = sh.Cm[3 * e]; elemF[b + 10] = sh.Cm[3 * e + 1]; elemF[b + 11] = sh.Cm[3 * e + 2];
55 elemF[b + 12] = 2 * setup.config.material.sigma1 * sh.rhoHA[e];
56 for (let i = 0; i < 81; i++) elemF[b + 13 + i] = sh.Ke[81 * e + i];
57 }
58 elemF.set(rad.taps, ELEM_STRIDE * E);
59 const elemU = new Uint32Array(3 * E); elemU.set(mesh.tri);
60 const nodeF = new Float32Array(NODE_STRIDE * N);
61 for (let n = 0; n < N; n++) {
62 const b = NODE_STRIDE * n;
63 for (let d = 0; d < 5; d++) { nodeF[b + d] = sh.Minv[5 * n + d]; nodeF[b + 5 + d] = sh.M[5 * n + d]; }
64 for (let k = 0; k < 6; k++) for (let d = 0; d < 5; d++) nodeF[b + 10 + 5 * k + d] = sh.Rm[(k * N + n) * 5 + d];
65 }
66 const adj = new Uint32Array(N + 1 + sh.adjEntries.length); adj.set(sh.adjOffsets); adj.set(sh.adjEntries, N + 1);
67 const hitsF = new Float32Array(HF * K), hitsU = new Uint32Array(HU * K);
68 for (let k = 0; k < K; k++) {
69 const hf = HF * k, hu = HU * k;
70 hitsF[hf] = sc.kPoint[k]; hitsF[hf + 1] = sc.alpha[k]; hitsF[hf + 2] = sc.mass[k]; hitsF[hf + 3] = sc.gap[k];
71 hitsF[hf + 4] = sc.velocity[k]; hitsF[hf + 5] = HOLD_HEIGHT; hitsF[hf + 6] = 0; hitsF[hf + 7] = 0; hitsF[hf + 8] = 0;
72 for (let p = 0; p < P; p++) {
73 for (let i = 0; i < 3; i++) { hitsF[hf + HIT_HEADER + 3 * p + i] = sc.weights[(k * P + p) * 3 + i]; hitsU[hu + HITU_HEADER + 3 * p + i] = sc.nodes[(k * P + p) * 3 + i]; }
74 hitsF[hf + HIT_HEADER + 3 * P + p] = sc.offsets[k * P + p];
75 }
76 hitsU[hu] = sc.releaseStep[k]; hitsU[hu + 1] = 0;
77 }
78 const radCoef = new Float32Array(6 * E);
79 for (let e = 0; e < E; e++) {
80 radCoef[6 * e] = rad.coefF[e]; radCoef[6 * e + 1] = rad.coefB[e]; radCoef[6 * e + 2] = rad.i0F[e];
81 radCoef[6 * e + 3] = rad.frF[e]; radCoef[6 * e + 4] = rad.i0B[e]; radCoef[6 * e + 5] = rad.frB[e];
82 }
83 const mk = (name: string, data: Float32Array | Uint32Array, usage = storageUsage()) => {
84 const buf = device.createBuffer({ label: name, size: Math.max(16, data.byteLength), usage });
85 device.queue.writeBuffer(buf, 0, data.buffer, data.byteOffset, data.byteLength);
86 this.buffers[name] = buf;
87 return buf;
88 };
89 mk("elemF", elemF); mk("elemU", elemU); mk("nodeF", nodeF); mk("adj", adj); mk("radCoef", radCoef);
90 mk("hitsF", hitsF); mk("hitsU", hitsU);
91 mk("state", new Float32Array(12 * N)); mk("elemForce", new Float32Array(15 * E));
92 mk("sync", new Float32Array(SYNC_HEADER + 2 * G * 12)); mk("rad", new Float32Array(E * W + E * D));
93 mk("pressure", new Float32Array(2 * chunkCapacity));
94 const driveF = new Float32Array(N);
95 if (opts.drive) for (let n = 0; n < N; n++) driveF[n] = opts.drive.fNode[n];
96 mk("driveF", driveF);
97 mk("phiM", opts.projection ? opts.projection.phiM : new Float32Array(0));
98 mk("projOut", new Float32Array(Math.max(1, this.projCapacity * 2 * this.nModes)));
99 mk("traceOut", new Float32Array(this.tracing ? chunkCapacity : 1));
100 const chunkBuf = device.createBuffer({ label: "chunk", size: CHUNK_UNIFORM_BYTES, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST });
101 this.buffers.chunk = chunkBuf;
102 const staging = (size: number) => device.createBuffer({ size, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST });
103 this.stagingPressure = staging(8 * chunkCapacity); this.stagingState = staging(12 * N * 4);
104 this.stagingHits = staging(Math.max(16, HF * K * 4)); this.stagingHitsU = staging(Math.max(16, HU * K * 4));
105 if (this.framesCapacity > 0) {
106 mk("movie", new Float32Array(this.framesCapacity * N));
107 this.stagingMovie = staging(this.framesCapacity * N * 4);
108 } else this.stagingMovie = null;
109 this.stagingProj = this.projCapacity > 0 ? staging(this.projCapacity * 2 * this.nModes * 4) : null;
110 this.stagingTrace = this.tracing ? staging(chunkCapacity * 4) : null;
112 // ---- pipelines
113 const drive = opts.drive;
114 const module = device.createShaderModule({ code: shaderSource({
115 N, E, K, P, M, W, D, G, dt, sigma0: setup.config.material.sigma0, nl: sh.nl,
116 Nm: Math.max(1, this.nModes),
117 driveOmega: drive ? drive.omega : 0, driveRampSteps: drive ? drive.rampSteps : 1, driveStopStep: drive ? drive.stopStep : 0,
118 traceMode: this.tracing ? opts.traceMode! : -1,
119 }) });
120 const bgl = (entries: [number, "read-only-storage" | "storage" | "uniform"][]) => device.createBindGroupLayout({
121 entries: entries.map(([binding, type]) => ({ binding, visibility: GPUShaderStage.COMPUTE, buffer: { type } })),
122 });
123 const L0 = bgl([[0, "read-only-storage"], [1, "read-only-storage"], [2, "read-only-storage"], [3, "storage"], [4, "storage"], [5, "read-only-storage"], [6, "storage"], [7, "storage"]]);
124 const L1 = bgl([[0, "storage"], [1, "read-only-storage"], [2, "read-only-storage"], [3, "read-only-storage"], [4, "storage"], [5, "storage"], [6, "storage"], [7, "read-only-storage"], [8, "uniform"]]);
125 const L2 = bgl([[0, "read-only-storage"], [1, "read-only-storage"], [2, "read-only-storage"], [3, "storage"], [4, "uniform"]]);
126 const L3 = bgl([[0, "read-only-storage"], [1, "storage"], [2, "uniform"], [3, "read-only-storage"]]);
127 const L4 = bgl([[0, "read-only-storage"], [2, "uniform"], [3, "read-only-storage"], [4, "read-only-storage"], [5, "storage"]]);
128 const L5 = bgl([[0, "read-only-storage"], [2, "uniform"], [3, "read-only-storage"], [4, "read-only-storage"], [6, "storage"]]);
129 const empty = device.createBindGroupLayout({ entries: [] });
130 const mkPipe = (entryPoint: string, layouts: GPUBindGroupLayout[]) => device.createComputePipeline({
131 label: entryPoint, layout: device.createPipelineLayout({ bindGroupLayouts: layouts }), compute: { module, entryPoint },
132 });
133 this.pipelines.element = mkPipe("elementPass", [L0]);
134 this.pipelines.node = mkPipe("nodePass", [empty, L1]);
135 this.pipelines.gather = mkPipe("gatherPass", [empty, empty, L2]);
136 if (this.framesCapacity > 0) this.pipelines.capture = mkPipe("capturePass", [empty, empty, empty, L3]);
137 if (this.projCapacity > 0) this.pipelines.proj = mkPipe("projPass", [empty, empty, empty, L4]);
138 if (this.tracing) this.pipelines.trace = mkPipe("tracePass", [empty, empty, empty, L5]);
139 const B = this.buffers;
140 const bg = (layout: GPUBindGroupLayout, entries: [number, GPUBuffer][]) => device.createBindGroup({ layout, entries: entries.map(([binding, buffer]) => ({ binding, resource: { buffer } })) });
141 this.bindGroups.element = bg(L0, [[0, B.elemF], [1, B.elemU], [2, B.state], [3, B.elemForce], [4, B.hitsF], [5, B.hitsU], [6, B.sync], [7, B.rad]]);
142 this.bindGroups.node = bg(L1, [[0, B.state], [1, B.elemForce], [2, B.nodeF], [3, B.adj], [4, B.hitsF], [5, B.hitsU], [6, B.sync], [7, B.driveF], [8, B.chunk]]);
143 this.bindGroups.gather = bg(L2, [[0, B.sync], [1, B.rad], [2, B.radCoef], [3, B.pressure], [4, B.chunk]]);
144 if (this.framesCapacity > 0) this.bindGroups.capture = bg(L3, [[0, B.state], [1, B.movie], [2, B.chunk], [3, B.sync]]);
145 if (this.projCapacity > 0) this.bindGroups.proj = bg(L4, [[0, B.state], [2, B.chunk], [3, B.sync], [4, B.phiM], [5, B.projOut]]);
146 if (this.tracing) this.bindGroups.trace = bg(L5, [[0, B.state], [2, B.chunk], [3, B.sync], [4, B.phiM], [6, B.traceOut]]);
147 this.bindGroups.empty = device.createBindGroup({ layout: empty, entries: [] });
148 }
150 setState(q: Float64Array, v: Float64Array) {
151 const N = this.setup.mesh.n, st = new Float32Array(12 * N);
152 for (let i = 0; i < 5 * N; i++) { st[i] = q[i]; st[5 * N + i] = v[i]; }
153 for (let n = 0; n < N; n++) st[11 * N + n] = v[5 * n + 2]; // vwPrev, so the first radiated acceleration is not a spike
154 this.device.queue.writeBuffer(this.buffers.state, 0, st);
155 }
157 /** Audio sample completed by the gather pass at sub-step n, or -1 when none is due. */
158 private sampleDueAt(n: number): number {
159 const M = this.setup.substeps, W = DECIMATION_SPAN;
160 if (n < 1 || (n - 1) % M !== 0) return -1;
161 const j0 = (n - 1) / M;
162 if (j0 === 0) return -1;
163 const j = j0 - W / 2;
164 return j >= 0 ? j : -1;
165 }
167 get maxChunk() { return this.chunkCapacity; }
169 async runChunk(nSamples: number): Promise<ChunkResult> {
170 if (nSamples > this.chunkCapacity) throw new Error("chunk larger than the pressure buffer");
171 const { device, setup } = this, M = setup.substeps, E = setup.mesh.e, K = setup.schedule.K, N = setup.mesh.n;
172 const n0 = this.step, n1 = n0 + nSamples * M;
173 let base = -1, count = 0;
174 for (let n = n0; n < n1; n++) { const j = this.sampleDueAt(n); if (j >= 0) { if (base < 0) base = j; count++; } }
175 // frames captured in this chunk: the state at the start of sub-step n where n / M is a multiple of movieEvery
176 const frameSamples: number[] = [];
177 const every = this.movieEvery;
178 if (every > 0) for (let n = n0; n < n1; n++) if (n % (M * every) === 0) frameSamples.push(n / M);
179 const frameBase = frameSamples.length ? frameSamples[0] / every : 0;
180 const projSamples: number[] = [];
181 const pEvery = this.projEvery;
182 if (pEvery > 0) for (let n = n0; n < n1; n++) if (n % (M * pEvery) === 0) projSamples.push(n / M);
183 const projBase = projSamples.length ? projSamples[0] / pEvery : 0;
184 // drive phase at the first sub-step of the chunk, computed in float64 and folded to [0, 2 pi)
185 const drv = this.opts.drive;
186 const phase0 = drv ? (drv.omega * setup.dt * n0) % (2 * Math.PI) : 0;
187 const u = new ArrayBuffer(CHUNK_UNIFORM_BYTES), dv = new DataView(u);
188 dv.setInt32(0, base < 0 ? 0 : base, true); dv.setUint32(4, Math.max(count, 1), true);
189 dv.setUint32(8, frameBase, true); dv.setUint32(12, Math.max(every, 1), true);
190 dv.setUint32(16, projBase, true); dv.setUint32(20, Math.max(pEvery, 1), true);
191 dv.setFloat32(24, phase0, true); dv.setUint32(28, n0, true);
192 device.queue.writeBuffer(this.buffers.chunk, 0, u);
193 const enc = device.createCommandEncoder();
194 enc.clearBuffer(this.buffers.pressure);
195 const pass = enc.beginComputePass();
196 const wgE = Math.ceil((E + K) / 64), wgN = this.G;
197 for (let n = n0; n < n1; n++) {
198 if (every > 0 && n % (M * every) === 0) {
199 pass.setPipeline(this.pipelines.capture); pass.setBindGroup(0, this.bindGroups.empty); pass.setBindGroup(1, this.bindGroups.empty);
200 pass.setBindGroup(2, this.bindGroups.empty); pass.setBindGroup(3, this.bindGroups.capture); pass.dispatchWorkgroups(wgN);
201 }
202 if (pEvery > 0 && n % (M * pEvery) === 0) {
203 pass.setPipeline(this.pipelines.proj); pass.setBindGroup(0, this.bindGroups.empty); pass.setBindGroup(1, this.bindGroups.empty);
204 pass.setBindGroup(2, this.bindGroups.empty); pass.setBindGroup(3, this.bindGroups.proj); pass.dispatchWorkgroups(this.nModes);
205 }
206 if (this.tracing && n % M === 0) {
207 pass.setPipeline(this.pipelines.trace); pass.setBindGroup(0, this.bindGroups.empty); pass.setBindGroup(1, this.bindGroups.empty);
208 pass.setBindGroup(2, this.bindGroups.empty); pass.setBindGroup(3, this.bindGroups.trace); pass.dispatchWorkgroups(1);
209 }
210 pass.setPipeline(this.pipelines.element); pass.setBindGroup(0, this.bindGroups.element); pass.dispatchWorkgroups(wgE);
211 if (this.sampleDueAt(n) >= 0) {
212 pass.setPipeline(this.pipelines.gather); pass.setBindGroup(0, this.bindGroups.empty); pass.setBindGroup(1, this.bindGroups.empty);
213 pass.setBindGroup(2, this.bindGroups.gather); pass.dispatchWorkgroups(1);
214 }
215 pass.setPipeline(this.pipelines.node); pass.setBindGroup(0, this.bindGroups.empty); pass.setBindGroup(1, this.bindGroups.node); pass.dispatchWorkgroups(wgN);
216 }
217 pass.end();
218 enc.copyBufferToBuffer(this.buffers.pressure, 0, this.stagingPressure, 0, 8 * this.chunkCapacity);
219 const nFrames = frameSamples.length;
220 if (nFrames > 0) enc.copyBufferToBuffer(this.buffers.movie, 0, this.stagingMovie!, 0, nFrames * N * 4);
221 const nProj = projSamples.length;
222 if (nProj > 0) enc.copyBufferToBuffer(this.buffers.projOut, 0, this.stagingProj!, 0, nProj * 2 * this.nModes * 4);
223 if (this.tracing) enc.copyBufferToBuffer(this.buffers.traceOut, 0, this.stagingTrace!, 0, nSamples * 4);
224 device.queue.submit([enc.finish()]);
225 this.step = n1; this.samplesDone += nSamples;
226 await this.stagingPressure.mapAsync(GPUMapMode.READ);
227 const all = new Float32Array(this.stagingPressure.getMappedRange().slice(0));
228 this.stagingPressure.unmap();
229 const frames: Float32Array[] = [];
230 if (nFrames > 0) {
231 await this.stagingMovie!.mapAsync(GPUMapMode.READ, 0, nFrames * N * 4);
232 const data = new Float32Array(this.stagingMovie!.getMappedRange(0, nFrames * N * 4).slice(0));
233 this.stagingMovie!.unmap();
234 for (let f = 0; f < nFrames; f++) frames.push(data.slice(f * N, (f + 1) * N));
235 }
236 const proj: Float32Array[] = [];
237 if (nProj > 0) {
238 const stride = 2 * this.nModes;
239 await this.stagingProj!.mapAsync(GPUMapMode.READ, 0, nProj * stride * 4);
240 const data = new Float32Array(this.stagingProj!.getMappedRange(0, nProj * stride * 4).slice(0));
241 this.stagingProj!.unmap();
242 for (let s = 0; s < nProj; s++) proj.push(data.slice(s * stride, (s + 1) * stride));
243 }
244 let trace = new Float32Array(0);
245 if (this.tracing) {
246 await this.stagingTrace!.mapAsync(GPUMapMode.READ, 0, nSamples * 4);
247 trace = new Float32Array(this.stagingTrace!.getMappedRange(0, nSamples * 4).slice(0));
248 this.stagingTrace!.unmap();
249 }
250 const size = Math.max(count, 1);
251 return { firstSample: base < 0 ? this.samplesDone : base, pressureFront: all.slice(0, count), pressureBack: all.slice(size, size + count),
252 frames, frameSamples, proj, projSamples, trace, traceBase: n0 / M };
253 }
255 private async read(buffer: GPUBuffer, staging: GPUBuffer, bytes: number): Promise<ArrayBuffer> {
256 const enc = this.device.createCommandEncoder();
257 enc.copyBufferToBuffer(buffer, 0, staging, 0, bytes);
258 this.device.queue.submit([enc.finish()]);
259 await staging.mapAsync(GPUMapMode.READ);
260 const out = staging.getMappedRange().slice(0);
261 staging.unmap();
262 return out;
263 }
265 async diagnostics(): Promise<Diagnostics> {
266 const { setup } = this, sh = setup.shell, N = setup.mesh.n, sc = setup.schedule, HF = this.HF, HU = this.HU;
267 const st = new Float32Array(await this.read(this.buffers.state, this.stagingState, 12 * N * 4));
268 const hf = new Float32Array(await this.read(this.buffers.hitsF, this.stagingHits, Math.max(16, HF * sc.K * 4)));
269 const hu = new Uint32Array(await this.read(this.buffers.hitsU, this.stagingHitsU, Math.max(16, HU * sc.K * 4)));
270 let ke = 0, woh = 0;
271 for (let i = 0; i < 5 * N; i++) ke += 0.5 * sh.M[i] * st[5 * N + i] * st[5 * N + i];
272 for (let n = 0; n < N; n++) woh = Math.max(woh, Math.abs(st[5 * n + 2] / sh.hNode[n]));
273 const { membrane, bending } = strainEnergies(sh, st.subarray(0, 5 * N));
274 let ex = 0;
275 const hits = [];
276 for (let k = 0; k < sc.K; k++) {
277 ex += 0.5 * hf[HF * k + 2] * hf[HF * k + 6] ** 2;
278 hits.push({ peakForce: hf[HF * k + 8], contactTime: hu[HU * k + 1] * setup.dt, velocity: hf[HF * k + 6], height: hf[HF * k + 5] });
279 }
280 return { time: this.step * setup.dt, kinetic: ke, membrane, bending, exciters: ex, maxWoverH: woh, hits };
281 }
283 async displacement(): Promise<Float32Array> {
284 const N = this.setup.mesh.n;
285 const st = new Float32Array(await this.read(this.buffers.state, this.stagingState, 12 * N * 4));
286 const w = new Float32Array(N);
287 for (let n = 0; n < N; n++) w[n] = st[5 * n + 2];
288 return w;
289 }
291 /** The state buffer (q at offset 0, 5 floats per node), for a renderer that reads w directly. */
292 get stateBuffer(): GPUBuffer { return this.buffers.state; }
294 dispose() {
295 for (const b of Object.values(this.buffers)) b.destroy();
296 this.stagingPressure.destroy(); this.stagingState.destroy(); this.stagingHits.destroy(); this.stagingHitsU.destroy();
297 this.stagingMovie?.destroy(); this.stagingProj?.destroy(); this.stagingTrace?.destroy();
298 }
↑↓ move↵openescclose