1import type { BeadSystem, Excitation, SimOptions, SimResult } from './types';
2import { timeStep, type Progress } from './solver';
4// The same scheme as simulate() in solver.ts (velocity Verlet, same time step,
5// rigid-motion removal, contact, delayed pickup), run on the GPU in single
6// precision. Each phase of a time step is its own small kernel, and a command
7// buffer holds the dispatches for many samples; time and sample counters live
8// on the GPU, so nothing returns to JavaScript until a chunk is done. (One
9// persistent workgroup running whole time steps was tried first: its large
10// kernel made every barrier cost several microseconds on an Intel iGPU, and it
11// could use only one of the GPU's subslices.) Forces are gathered per bead from
12// fixed-width adjacency lists, so no atomics are needed. Sums over beads go
13// through per-workgroup partials, which each workgroup of the next kernel adds
14// up itself (cheaper than a dispatch of its own). Three buffers only: the
15// per-dispatch cost grows with the number bound. To keep small vibrations from vanishing in
16// f32, the state is the displacement from rest, and spring extensions and
17// hinge angles are built from precomputed rest differences.
19const WGB = 64; // invocations per workgroup
21interface Packed {
22 wgsl: string;
23 ro: Uint32Array; // parameters (f32 bits) and indices
24 rw: Float32Array; // state, shared scalars, per-workgroup partial sums
25 outLen: number; // audio, contact force, frame scalars, mallet positions, then frames
26 framesAt: number;
27 nSamples: number;
28 nFrames: number;
29 samplesPerFrame: number;
30 dt: number;
31 stepsPerSample: number;
32 omegaMax: number;
33 nBead: number; // workgroups over beads
34 nHinge: number; // workgroups over hinges
35 hasMallet: boolean;
36 freeRigid: boolean;
37 velForces: boolean;
38}
40function f(x: number): string {
41 if (!Number.isFinite(x)) throw new Error(`non-finite constant ${x}`);
42 return x.toExponential(9);
43}
45class Pack<T extends Float32Array | Uint32Array> {
46 parts: ArrayLike<number>[] = [];
47 len = 0;
48 constructor(private Ctor: new (n: number) => T) {}
49 add(a: ArrayLike<number>): number {
50 const off = this.len;
51 this.parts.push(a);
52 this.len += a.length;
53 return off;
54 }
55 build(): T {
56 const out = new this.Ctor(Math.max(4, this.len));
57 let o = 0;
58 for (const p of this.parts) {
59 out.set(p, o);
60 o += p.length;
61 }
62 return out;
63 }
64}
66// rows of variable length -> slot-major table [slot][row] padded with `pad`
67function padded<T>(rows: T[][], pad: T): { width: number; table: T[] } {
68 const width = Math.max(1, ...rows.map((r) => r.length));
69 const table: T[] = [];
70 for (let q = 0; q < width; q++) for (const r of rows) table.push(q < r.length ? r[q] : pad);
71 return { width, table };
72}
74function pack(sys: BeadSystem, exc: Excitation, opt: SimOptions, sg: boolean): Packed {
75 const { n, si, sj, sk, sL0, sc, mass, pinned, anchors, rest } = sys;
76 const mallet = exc.mallet;
77 const { omegaMax, dt, stepsPerSample } = timeStep(sys, mallet, opt.fs, opt.cfl);
78 const nSamples = Math.round(opt.duration * opt.fs);
79 const samplesPerFrame = Math.max(1, Math.round(opt.fs / opt.frameRate));
80 const nFrames = Math.floor((nSamples - 1) / samplesPerFrame) + 1;
82 const prm = new Pack(Float32Array);
83 const idx = new Pack(Uint32Array);
85 // beads
86 const invMass = Float64Array.from(mass, (m, i) => (pinned[i] ? 0 : 1 / m));
87 let M = 0;
88 const cm = [0, 0, 0];
89 for (let i = 0; i < n; i++) {
90 M += mass[i];
91 for (let c = 0; c < 3; c++) cm[c] += mass[i] * rest[3 * i + c];
92 }
93 const r0 = new Float64Array(3 * n);
94 for (let i = 0; i < n; i++) for (let c = 0; c < 3; c++) r0[3 * i + c] = rest[3 * i + c] - cm[c] / M;
95 const i0 = new Uint32Array(n);
96 const fr = new Float64Array(n);
97 if (opt.pickupDelay)
98 for (let i = 0; i < n; i++) {
99 i0[i] = Math.floor(opt.pickupDelay[i]);
100 fr[i] = opt.pickupDelay[i] - i0[i];
101 }
102 const H = Math.max(...i0) + 2;
103 const pw = mallet ? mallet.weights : new Float64Array(n);
104 const patch: number[] = [];
105 for (let i = 0; i < n; i++) if (pw[i] > 0) patch.push(i);
106 let mEff = 1;
107 if (mallet) {
108 let inv = 0;
109 for (const i of patch) inv += (pw[i] * pw[i]) / mass[i];
110 mEff = 1 / (1 / mallet.mass + inv);
111 }
113 const P_MASS = prm.add(mass);
114 const P_IM = prm.add(invMass);
115 const P_R0 = prm.add(r0);
116 const P_W = prm.add(opt.pickup);
117 const P_FR = prm.add(fr);
118 const P_PW = prm.add(pw);
119 const P_FS = prm.add(exc.force ? exc.force.shape : new Float64Array(3 * n));
121 // springs, gathered per bead: [k, c, rest difference (3), rest length, e0, weight]
122 // e0 = rest length - natural length (pre-tension); the weight counts a spring
123 // between two moving beads half at each end. Padding: a zero spring to a fixed point.
124 const FIXED = 0xffffffff;
125 const SLIDE = 0x80000000;
126 const code = (node: number) => (node < n ? node : anchors[node - n].slideBead < 0 ? FIXED : SLIDE + anchors[node - n].slideBead);
127 const nodePos = (node: number, c: number) => {
128 if (node < n) return rest[3 * node + c];
129 const a = anchors[node - n];
130 return a.slideBead < 0 || c === 0 ? a.pos[c] : rest[3 * a.slideBead + c];
131 };
132 type SpringEntry = { other: number; vals: number[] };
133 const byBead: SpringEntry[][] = Array.from({ length: n }, () => []);
134 for (let s = 0; s < si.length; s++) {
135 for (const [me, ot] of [[si[s], sj[s]], [sj[s], si[s]]]) {
136 if (me >= n || pinned[me]) continue;
137 const rd = [0, 1, 2].map((c) => nodePos(ot, c) - nodePos(me, c));
138 const lr = Math.hypot(rd[0], rd[1], rd[2]);
139 const moving = ot < n && !pinned[ot];
140 byBead[me].push({ other: code(ot), vals: [sk[s], sc[s], rd[0], rd[1], rd[2], lr, lr - sL0[s], moving ? 0.5 : 1] });
141 }
142 }
143 const sp = padded(byBead, { other: FIXED, vals: [0, 0, 1, 0, 0, 1, 0, 0] });
144 const P_SF = prm.add(sp.table.flatMap((e) => e.vals));
145 const I_SO = idx.add(sp.table.map((e) => e.other));
147 // bending: hinge nh is a dummy with zero angle, rate, moment and gradient
148 const b = sys.bending;
149 const nh = b ? b.w0.length : 0;
150 let P_HF = 0, P_HC = 0, I_HN = 0, I_HCC = 0, I_BHE = 0;
151 let MAXHC = 1, MAXBH = 1;
152 if (b) {
153 const hf: number[] = [];
154 const hn: number[] = [];
155 const d = (p: number, q: number) => [0, 1, 2].map((c) => rest[3 * p + c] - rest[3 * q + c]);
156 for (let e = 0; e < nh; e++) {
157 const [w0, w1, e0, e1] = [b.w0[e], b.w1[e], b.e0[e], b.e1[e]];
158 if (Math.max(w0, w1, e0, e1) >= n) throw new Error('hinge on an anchor');
159 hn.push(w0, w1, e0, e1);
160 hf.push(b.sign[e], b.theta0[e], ...d(e1, e0), ...d(w0, e0), ...d(w1, e1));
161 }
162 // moment of hinge a = sum over its triangles of M[a][c] phi_c
163 const rows: Map<number, number>[] = Array.from({ length: nh }, () => new Map());
164 for (let t = 0; t < b.tri.length / 3; t++)
165 for (let a = 0; a < 3; a++) {
166 const ea = b.tri[3 * t + a];
167 if (ea < 0) continue;
168 for (let c = 0; c < 3; c++) {
169 const ec = b.tri[3 * t + c];
170 if (ec >= 0) rows[ea].set(ec, (rows[ea].get(ec) ?? 0) + b.M[9 * t + 3 * a + c]);
171 }
172 }
173 const hc = padded(rows.map((r) => [...r]), [nh, 0]);
174 MAXHC = hc.width;
175 // per bead: the hinges it belongs to (hinge * 4 + slot)
176 const bh: number[][] = Array.from({ length: n }, () => []);
177 for (let e = 0; e < nh; e++) for (let k = 0; k < 4; k++) if (!pinned[hn[4 * e + k]]) bh[hn[4 * e + k]].push(4 * e + k);
178 const bhp = padded(bh, 4 * nh);
179 MAXBH = bhp.width;
180 P_HF = prm.add(hf);
181 P_HC = prm.add(hc.table.map(([, v]) => v));
182 I_HN = idx.add(hn);
183 I_HCC = idx.add(hc.table.map(([c]) => c));
184 I_BHE = idx.add(bhp.table);
185 }
186 const I_PATCH = idx.add(patch);
187 const I_I0 = idx.add(i0);
189 // state
190 const S_U = 0;
191 const S_V = S_U + 3 * n;
192 const S_A = S_V + 3 * n;
193 const S_G = S_A + 3 * n;
194 const S_PHI = S_G + 12 * (nh + 1);
195 const S_RATE = S_PHI + nh + 1;
196 const S_MOM = S_RATE + nh + 1;
197 const S_MEL = S_MOM + nh + 1;
198 const S_HIST = S_MEL + nh + 1;
199 const state = new Float32Array(S_HIST + n * H);
200 for (let i = 0; i < n; i++)
201 if (!pinned[i])
202 for (let c = 0; c < 3; c++) {
203 state[S_U + 3 * i + c] = exc.u0[3 * i + c];
204 state[S_V + 3 * i + c] = exc.v0[3 * i + c];
205 }
206 const nBead = Math.ceil(n / WGB);
207 const nHinge = Math.ceil(nh / WGB);
208 // read-write buffer: the state, then shared scalars, then per-workgroup partial sums
209 const SC = state.length;
210 const Q_BF = SC + 64; // 17 per bead workgroup
211 const Q_HD = Q_BF + 17 * nBead; // 1 per hinge workgroup (hinge damping power)
212 const Q_KC = Q_HD + Math.max(1, nHinge); // 6 per bead workgroup
213 const Q_REC = Q_KC + 6 * nBead; // 3 per bead or hinge workgroup: audio, energy, largest strain
214 const Q_IRV = Q_REC + 3 * (nBead + nHinge); // 15 per bead workgroup
215 const rw = new Float32Array(Q_IRV + 15 * nBead);
216 rw.set(state);
217 if (mallet) rw.set(mallet.vel0, SC + 3);
218 // read-only buffer: parameters (as f32 bits), then indices
219 const prmArr = prm.build();
220 const idxArr = idx.build();
221 const ro = new Uint32Array(prmArr.length + idxArr.length);
222 ro.set(new Uint32Array(prmArr.buffer));
223 ro.set(idxArr, prmArr.length);
224 const IO = prmArr.length;
226 const O_AUD = 0;
227 const O_CF = nSamples;
228 const O_FSC = 2 * nSamples;
229 const O_MP = O_FSC + 6 * nFrames;
230 const OFR = O_MP + 3 * nFrames;
232 const force = exc.force;
233 const sigma1 = b?.sigma1 ?? 0;
234 const anyDash = sc.some((c) => c !== 0);
235 const velForces = sys.sigma0 !== 0 || anyDash || sigma1 !== 0 || (mallet?.mu ?? 0) > 0;
236 const u = (x: number) => `${x}u`;
237 const consts = `
238const N = ${u(n)};
239const NH = ${u(nh)};
240const NP = ${u(patch.length)};
241const NBW = ${u(nBead)};
242const NHW = ${u(nHinge)};
243const MAXS = ${u(sp.width)};
244const MAXBH = ${u(MAXBH)};
245const MAXHC = ${u(MAXHC)};
246const H = ${u(H)};
247const SPF = ${u(samplesPerFrame)};
248const NF = ${u(nFrames)};
249const DT = ${f(dt)};
250const S0 = ${f(sys.sigma0)};
251const S1 = ${f(sigma1)};
252const MTOT = ${f(M)};
253const HAS_MALLET = ${!!mallet};
254const FREE_RIGID = ${!!sys.freeRigid};
255const HAS_PULSE = ${!!force};
256const ANY_DASH = ${anyDash};
257const MM = ${f(mallet?.mass ?? 1)};
258const MK = ${f(mallet?.K ?? 0)};
259const MP = ${f(mallet?.p ?? 1)};
260const MU = ${f(mallet?.mu ?? 0)};
261const REACH2 = ${f(mallet?.reach !== undefined ? mallet.reach * mallet.reach : -1)};
262const MEFF = ${f(mEff)};
263const NRM = vec3f(${(mallet?.normal ?? [0, 0, 0]).map(f).join(', ')});
264const F0 = ${f(force?.F0 ?? 0)};
265const TAU = ${f(force?.tau ?? 1)};
266const P_MASS = ${u(P_MASS)}; const P_IM = ${u(P_IM)}; const P_R0 = ${u(P_R0)}; const P_W = ${u(P_W)};
267const P_FR = ${u(P_FR)}; const P_PW = ${u(P_PW)}; const P_FS = ${u(P_FS)}; const P_SF = ${u(P_SF)};
268const P_HF = ${u(P_HF)}; const P_HC = ${u(P_HC)};
269const I_SO = ${u(I_SO + IO)}; const I_HN = ${u(I_HN + IO)}; const I_HCC = ${u(I_HCC + IO)}; const I_BHE = ${u(I_BHE + IO)};
270const I_PATCH = ${u(I_PATCH + IO)}; const I_I0 = ${u(I_I0 + IO)};
271const S_U = ${u(S_U)}; const S_V = ${u(S_V)}; const S_A = ${u(S_A)}; const S_G = ${u(S_G)};
272const S_PHI = ${u(S_PHI)}; const S_RATE = ${u(S_RATE)}; const S_MOM = ${u(S_MOM)}; const S_MEL = ${u(S_MEL)};
273const S_HIST = ${u(S_HIST)}; const SC = ${u(SC)};
274const Q_BF = ${u(Q_BF)}; const Q_HD = ${u(Q_HD)}; const Q_KC = ${u(Q_KC)}; const Q_REC = ${u(Q_REC)}; const Q_IRV = ${u(Q_IRV)};
275const O_AUD = ${u(O_AUD)}; const O_CF = ${u(O_CF)}; const O_FSC = ${u(O_FSC)}; const O_MP = ${u(O_MP)}; const OFR = ${u(OFR)};
276`;
277 return {
278 wgsl: (sg ? 'enable subgroups;\n' : '') + consts + reducers(sg) + KERNELS.replaceAll('/*SG*/', sg ? ', @builtin(subgroup_invocation_id) sgiIn: u32, @builtin(subgroup_size) sgsIn: u32' : '').replaceAll('/*SGI*/', sg ? 'sgi = sgiIn; sgs = sgsIn;' : ''),
279 ro,
280 rw,
281 outLen: OFR + 3 * n * nFrames,
282 framesAt: OFR,
283 nSamples,
284 nFrames,
285 samplesPerFrame,
286 dt,
287 stepsPerSample,
288 omegaMax,
289 nBead,
290 nHinge,
291 hasMallet: !!mallet,
292 freeRigid: !!sys.freeRigid,
293 velForces,
294 };
295}
297// Workgroup sums: each invocation puts its values, then wsum(K) leaves the
298// sums in wres[0..K] (all invocations must call it). wmax: the largest value
299// put at 0. With subgroups: a subgroup sum per value, then one invocation per
300// value adds the subgroups' sums (subgroups taken to be consecutive
301// invocations); else a tree in workgroup memory.
302function reducers(sg: boolean): string {
303 const common = `
304var<workgroup> wred: array<f32, 17u * 64u>;
305var<workgroup> wres: array<f32, 17>;
306var<private> sgi: u32;
307var<private> sgs: u32;
308fn put(l: u32, k: u32, v: f32) { wred[k * 64u + l] = v; }
309fn put3(l: u32, k: u32, v: vec3f) { put(l, k, v.x); put(l, k + 1u, v.y); put(l, k + 2u, v.z); }
310`;
311 if (sg)
312 return common + `
313var<workgroup> wsg: array<f32, 17u * 16u>;
314fn wsum(l: u32, K: u32) {
315 workgroupBarrier();
316 let g = l / sgs;
317 for (var k = 0u; k < K; k++) {
318 let s = subgroupAdd(wred[k * 64u + l]);
319 if (sgi == 0u) { wsg[k * 16u + g] = s; }
320 }
321 workgroupBarrier();
322 if (l < K) {
323 var t = 0.0;
324 for (var j = 0u; j < 64u / sgs; j++) { t += wsg[l * 16u + j]; }
325 wres[l] = t;
326 }
327 workgroupBarrier();
328}
329fn wmax(l: u32) -> f32 {
330 workgroupBarrier();
331 let s = subgroupMax(wred[l]);
332 if (sgi == 0u) { wsg[l / sgs] = s; }
333 workgroupBarrier();
334 var m = 0.0;
335 for (var j = 0u; j < 64u / sgs; j++) { m = max(m, wsg[j]); }
336 workgroupBarrier();
337 return m;
338}
339`;
340 return common + `
341fn wsum(l: u32, K: u32) {
342 workgroupBarrier();
343 // each level adds pairs (k, j) and (k, j + s), spread over all invocations
344 for (var s = 32u; s > 0u; s = s / 2u) {
345 for (var x = l; x < K * s; x += 64u) {
346 let k = x / s;
347 let j = x - k * s;
348 wred[k * 64u + j] += wred[k * 64u + j + s];
349 }
350 workgroupBarrier();
351 }
352 if (l < K) { wres[l] = wred[l * 64u]; }
353 workgroupBarrier();
354}
355fn wmax(l: u32) -> f32 {
356 workgroupBarrier();
357 for (var s = 32u; s > 0u; s = s / 2u) {
358 if (l < s) { wred[l] = max(wred[l], wred[l + s]); }
359 workgroupBarrier();
360 }
361 let m = wred[0];
362 workgroupBarrier();
363 return m;
364}
365`;
366}
368// Shared scalars (at SC), written by invocation 0 of workgroup 0:
369// 0-2 mallet position q, 3-5 velocity, 6-8 acceleration, 9-11 force on the
370// mallet, 12 contact force, 13 contact energy, 14 friction power,
371// 17-18 dissipated since the last frame (Kahan sum), 19-20 work in since the
372// last frame, 44 time steps taken and 45 next sample to record (u32 bits).
373// Each workgroup keeps its own copy of the rigid-motion correction (wcor).
374const KERNELS = /* wgsl */ `
375@group(0) @binding(0) var<storage, read> ro: array<u32>;
376@group(0) @binding(1) var<storage, read_write> rw: array<f32>;
377@group(0) @binding(2) var<storage, read_write> ob: array<f32>;
379fn P(k: u32) -> f32 { return bitcast<f32>(ro[k]); }
380fn pv(o: u32) -> vec3f { return vec3f(P(o), P(o + 1u), P(o + 2u)); }
381fn ld(o: u32) -> vec3f { return vec3f(rw[o], rw[o + 1u], rw[o + 2u]); }
382fn sv(o: u32, x: vec3f) { rw[o] = x.x; rw[o + 1u] = x.y; rw[o + 2u] = x.z; }
383fn s3(k: u32) -> vec3f { return ld(SC + k); }
384fn ss3(k: u32, x: vec3f) { sv(SC + k, x); }
385fn count(k: u32) -> u32 { return bitcast<u32>(rw[SC + k]); }
386fn gU(i: u32) -> vec3f { return ld(S_U + 3u * i); }
387fn gV(i: u32) -> vec3f { return ld(S_V + 3u * i); }
388fn gA(i: u32) -> vec3f { return ld(S_A + 3u * i); }
389fn sU(i: u32, x: vec3f) { sv(S_U + 3u * i, x); }
390fn sV(i: u32, x: vec3f) { sv(S_V + 3u * i, x); }
391fn sA(i: u32, x: vec3f) { sv(S_A + 3u * i, x); }
392// position of bead i relative to the rest centre of mass
393fn pos(i: u32) -> vec3f { return pv(P_R0 + 3u * i) + gU(i); }
395// displacement or velocity of a spring's far node: a bead, a fixed anchor (0)
396// or a sliding anchor that follows bead b across its rod
397fn otherU(o: u32) -> vec3f {
398 if (o == 0xffffffffu) { return vec3f(0.0); }
399 if (o >= 0x80000000u) { let x = gU(o - 0x80000000u); return vec3f(0.0, x.y, x.z); }
400 return gU(o);
401}
402fn otherV(o: u32) -> vec3f {
403 if (o == 0xffffffffu) { return vec3f(0.0); }
404 if (o >= 0x80000000u) { let x = gV(o - 0x80000000u); return vec3f(0.0, x.y, x.z); }
405 return gV(o);
406}
408fn kahan(k: u32, x: f32) {
409 let o = SC + k;
410 let y = x - rw[o + 1u];
411 let t = rw[o] + y;
412 rw[o + 1u] = (t - rw[o]) - y;
413 rw[o] = t;
414}
416// sums of per-workgroup partials rw[q + w * stride + k] over w < nw into wres[0..K]
417// (the loads are spread over the invocations: global memory latency, not
418// arithmetic, is what costs here)
419fn psum(l: u32, q: u32, K: u32, stride: u32, nw: u32) {
420 for (var k = 0u; k < K; k++) { put(l, k, 0.0); }
421 for (var x = l; x < nw * K; x += 64u) {
422 let w = x / K;
423 let k = x - w * K;
424 wred[k * 64u + l] += rw[q + w * stride + k];
425 }
426 wsum(l, K);
427}
429// sums for the rigid part of a per-bead field x: mass-weighted geometry
430// (m u, m s s^T as diagonal and off-diagonal) and momentum (m x, m s x x),
431// s = position relative to the rest centre of mass
432struct Rig { mu: vec3f, sd: vec3f, so: vec3f, p: vec3f, l: vec3f };
433fn rigAdd(r: Rig, m: f32, u: vec3f, s: vec3f, x: vec3f) -> Rig {
434 return Rig(r.mu + m * u, r.sd + m * s * s, r.so + m * vec3f(s.x * s.y, s.x * s.z, s.y * s.z), r.p + m * x, r.l + m * cross(s, x));
435}
436fn rigZero() -> Rig { return Rig(vec3f(0.0), vec3f(0.0), vec3f(0.0), vec3f(0.0), vec3f(0.0)); }
437fn rigPut(l: u32, r: Rig) { put3(l, 0u, r.mu); put3(l, 3u, r.sd); put3(l, 6u, r.so); put3(l, 9u, r.p); put3(l, 12u, r.l); }
439// This workgroup's rigid-motion correction: 0-2 centre of mass, 3-8 inverse
440// inertia, 9-11 rigid translation rate, 12-14 rotation rate, 15 kinetic
441// energy of the rigid part, 16 damping power, 17 external power (from the
442// bead-force sums).
443var<workgroup> wcor: array<f32, 18>;
444fn w3(k: u32) -> vec3f { return vec3f(wcor[k], wcor[k + 1u], wcor[k + 2u]); }
445fn sw3(k: u32, x: vec3f) { wcor[k] = x.x; wcor[k + 1u] = x.y; wcor[k + 2u] = x.z; }
446// invocation 0: centre of mass and inverse inertia from geometry sums in wres[0..9]
447fn rigidGeometry() {
448 let c = vec3f(wres[0], wres[1], wres[2]) / MTOT;
449 let sxx = wres[3] - MTOT * c.x * c.x;
450 let syy = wres[4] - MTOT * c.y * c.y;
451 let szz = wres[5] - MTOT * c.z * c.z;
452 let sxy = wres[6] - MTOT * c.x * c.y;
453 let sxz = wres[7] - MTOT * c.x * c.z;
454 let syz = wres[8] - MTOT * c.y * c.z;
455 let a = syy + szz; let b = -sxy; let cc = -sxz;
456 let d = sxx + szz; let e = -syz; let ff = sxx + syy;
457 let A00 = d * ff - e * e; let A01 = cc * e - b * ff; let A02 = b * e - cc * d;
458 let A11 = a * ff - cc * cc; let A12 = b * cc - a * e; let A22 = a * d - b * b;
459 let det = a * A00 + b * A01 + cc * A02;
460 let k = select(0.0, 1.0 / det, abs(det) > 0.0);
461 wcor[3] = A00 * k; wcor[4] = A01 * k; wcor[5] = A02 * k;
462 wcor[6] = A11 * k; wcor[7] = A12 * k; wcor[8] = A22 * k;
463 sw3(0u, c);
464}
465// invocation 0: rigid translation and rotation rate of a field from its
466// momentum sums in wres[o..o+6], and its kinetic energy
467fn rigidPart(o: u32) {
468 let P = vec3f(wres[o], wres[o + 1u], wres[o + 2u]);
469 let L = vec3f(wres[o + 3u], wres[o + 4u], wres[o + 5u]) - cross(w3(0u), P);
470 let w = vec3f(wcor[3] * L.x + wcor[4] * L.y + wcor[5] * L.z, wcor[4] * L.x + wcor[6] * L.y + wcor[7] * L.z, wcor[5] * L.x + wcor[7] * L.y + wcor[8] * L.z);
471 let vc = P / MTOT;
472 sw3(9u, vc);
473 sw3(12u, w);
474 wcor[15] = 0.5 * dot(P, vc) + 0.5 * dot(w, L);
475}
476fn rigidAt(s: vec3f) -> vec3f { return w3(9u) + cross(w3(12u), s - w3(0u)); }
478// all invocations: the geometry and the rigid part of the accelerations from
479// the bead-force partials, and the damping and external power
480fn accelRigid(l: u32) {
481 psum(l, Q_BF, 17u, 17u, NBW);
482 if (l == 0u) {
483 rigidGeometry();
484 rigidPart(9u);
485 wcor[16] = wres[15];
486 wcor[17] = wres[16];
487 }
488 workgroupBarrier();
489}
490// all invocations: the geometry from the bead-force partials and the rigid
491// part of the velocities from the kick partials
492fn velRigid(l: u32) {
493 psum(l, Q_BF, 9u, 17u, NBW);
494 if (l == 0u) { rigidGeometry(); }
495 workgroupBarrier();
496 psum(l, Q_KC, 6u, 6u, NBW);
497 if (l == 0u) { rigidPart(0u); }
498 workgroupBarrier();
499}
501// dihedral angle change and its gradient (Bridson et al. 2003; see hinge.ts)
502fn hinge(e: u32) {
503 let w0 = ro[I_HN + 4u * e];
504 let w1 = ro[I_HN + 4u * e + 1u];
505 let e0 = ro[I_HN + 4u * e + 2u];
506 let e1 = ro[I_HN + 4u * e + 3u];
507 let hb = P_HF + 11u * e;
508 let sg = P(hb);
509 let ue0 = gU(e0);
510 let ue1 = gU(e1);
511 let E = pv(hb + 2u) + (ue1 - ue0);
512 let P1 = pv(hb + 5u) + (gU(w0) - ue0);
513 let P2 = pv(hb + 8u) + (gU(w1) - ue1);
514 let Q1 = P1 - E;
515 let Q2 = P2 + E;
516 let N1 = cross(P1, Q1);
517 let N2 = cross(P2, Q2);
518 let n1sq = dot(N1, N1);
519 let n2sq = dot(N2, N2);
520 let eLen = length(E);
521 let nn = sqrt(n1sq) * sqrt(n2sq);
522 let theta = atan2(dot(cross(N1, N2), E) / (nn * eLen), dot(N1, N2) / nn);
523 let gw0 = (-eLen / n1sq) * N1;
524 let gw1 = (-eLen / n2sq) * N2;
525 let t1 = dot(Q1, E) / eLen / n1sq;
526 let t2 = dot(P2, E) / eLen / n2sq;
527 let r1 = dot(P1, E) / eLen / n1sq;
528 let r2 = dot(Q2, E) / eLen / n2sq;
529 let ge0 = -(t1 * N1 + t2 * N2);
530 let ge1 = r1 * N1 + r2 * N2;
531 let g = S_G + 12u * e;
532 sv(g, sg * gw0);
533 sv(g + 3u, sg * gw1);
534 sv(g + 6u, sg * ge0);
535 sv(g + 9u, sg * ge1);
536 rw[S_PHI + e] = sg * (theta - P(hb + 1u));
537 if (S1 != 0.0) {
538 rw[S_RATE + e] = sg * (dot(gw0, gV(w0)) + dot(gw1, gV(w1)) + dot(ge0, gV(e0)) + dot(ge1, gV(e1)));
539 }
540}
542// mallet contact force from the patch-weighted displacement U and velocity Vp
543fn contact(U: vec3f, Vp: vec3f) {
544 let q = s3(0u);
545 let qv = s3(3u);
546 var fm = vec3f(0.0);
547 var fc = 0.0;
548 var ec = 0.0;
549 var pf = 0.0;
550 let d = dot(q - U, NRM);
551 if (d > 0.0) {
552 let o = q - U - d * NRM;
553 if (REACH2 < 0.0 || dot(o, o) <= REACH2) {
554 let fnorm = MK * pow(d, MP);
555 fc = fnorm;
556 ec = fnorm * d / (MP + 1.0);
557 fm = -fnorm * NRM;
558 if (MU > 0.0) {
559 var r = qv - Vp;
560 r -= dot(r, NRM) * NRM;
561 let slip = length(r);
562 let eps = max(1e-3, 4.0 * MU * fnorm * DT / MEFF);
563 let g = MU * fnorm / sqrt(slip * slip + eps * eps);
564 fm -= g * r;
565 pf = g * slip * slip;
566 }
567 }
568 }
569 ss3(6u, fm / MM);
570 ss3(9u, fm);
571 rw[SC + 12u] = fc;
572 rw[SC + 13u] = ec;
573 rw[SC + 14u] = pf;
574}
576// invocation 0 of workgroup 0 (after the mallet's position and velocity
577// are final for the step): the mallet's half kick and drift, count the step
578fn malletDrift() {
579 if (HAS_MALLET) {
580 let qv = s3(3u) + 0.5 * DT * s3(6u);
581 ss3(3u, qv);
582 ss3(0u, s3(0u) + DT * qv);
583 }
584 rw[SC + 44u] = bitcast<f32>(count(44u) + 1u);
585}
587// workgroup 0: sum the record partials and write the sample (and frame),
588// then move to the next sample
589fn recordSums(l: u32) {
590 var mx = 0.0;
591 for (var w = l; w < NBW + NHW; w += 64u) { mx = max(mx, rw[Q_REC + 3u * w + 2u]); }
592 put(l, 0u, mx);
593 let largest = wmax(l);
594 psum(l, Q_REC, 3u, 3u, NBW + NHW);
595 if (l == 0u) {
596 let smp = count(45u);
597 ob[O_AUD + smp] = wres[0];
598 ob[O_CF + smp] = rw[SC + 12u];
599 if (smp % SPF == 0u && smp / SPF < NF) {
600 let fi = smp / SPF;
601 let b = O_FSC + 6u * fi;
602 let qv = s3(3u);
603 ob[b] = wres[1];
604 ob[b + 1u] = select(0.0, 0.5 * MM * dot(qv, qv) + rw[SC + 13u], HAS_MALLET);
605 ob[b + 2u] = rw[SC + 17u];
606 ob[b + 3u] = rw[SC + 19u];
607 ob[b + 4u] = largest;
608 rw[SC + 17u] = 0.0; rw[SC + 18u] = 0.0; rw[SC + 19u] = 0.0; rw[SC + 20u] = 0.0;
609 let q = s3(0u);
610 ob[O_MP + 3u * fi] = q.x;
611 ob[O_MP + 3u * fi + 1u] = q.y;
612 ob[O_MP + 3u * fi + 2u] = q.z;
613 }
614 rw[SC + 45u] = bitcast<f32>(smp + 1u);
615 }
616}
618// first kick and drift of a sample's steps; workgroup 0 also writes the
619// sample recorded just before
620@compute @workgroup_size(64)
621fn kickDriftRec(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
622 /*SGI*/
623 if (wg.x == 0u) {
624 recordSums(l);
625 if (l == 0u) { malletDrift(); }
626 }
627 let i = wg.x * 64u + l;
628 if (i >= N || P(P_IM + i) == 0.0) { return; }
629 let v = gV(i) + 0.5 * DT * gA(i);
630 sV(i, v);
631 sU(i, gU(i) + DT * v);
632}
634// later kicks and drifts: first remove the rigid part of the velocities left
635// by the previous step
636@compute @workgroup_size(64)
637fn kickDrift(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
638 /*SGI*/
639 if (FREE_RIGID) { velRigid(l); }
640 if (wg.x == 0u && l == 0u) {
641 if (FREE_RIGID) { kahan(17u, wcor[15]); }
642 malletDrift();
643 }
644 let i = wg.x * 64u + l;
645 if (i >= N || P(P_IM + i) == 0.0) { return; }
646 var v = gV(i);
647 if (FREE_RIGID) { v -= rigidAt(pos(i)); }
648 v += 0.5 * DT * gA(i);
649 sV(i, v);
650 sU(i, gU(i) + DT * v);
651}
653// after the last step of a sample: remove the rigid part of the velocities
654@compute @workgroup_size(64)
655fn velCorrect(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
656 /*SGI*/
657 velRigid(l);
658 if (wg.x == 0u && l == 0u) { kahan(17u, wcor[15]); }
659 let i = wg.x * 64u + l;
660 if (i >= N) { return; }
661 sV(i, gV(i) - rigidAt(pos(i)));
662}
664// hinge angles, gradients and rates; the last workgroup does the mallet contact
665@compute @workgroup_size(64)
666fn hinges(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
667 /*SGI*/
668 if (wg.x < NHW) {
669 let e = wg.x * 64u + l;
670 if (e < NH) { hinge(e); }
671 return;
672 }
673 var su = vec3f(0.0);
674 var sv = vec3f(0.0);
675 for (var j = l; j < NP; j += 64u) {
676 let i = ro[I_PATCH + j];
677 let w = P(P_PW + i);
678 su += w * gU(i);
679 sv += w * gV(i);
680 }
681 put3(l, 0u, su);
682 put3(l, 3u, sv);
683 wsum(l, 6u);
684 if (l == 0u) { contact(vec3f(wres[0], wres[1], wres[2]), vec3f(wres[3], wres[4], wres[5])); }
685}
687// hinge moments, elastic and with damping; per-workgroup hinge damping power
688@compute @workgroup_size(64)
689fn moments(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
690 /*SGI*/
691 let e = wg.x * 64u + l;
692 var pd = 0.0;
693 if (e < NH) {
694 var m = 0.0;
695 var r = 0.0;
696 for (var q = 0u; q < MAXHC; q++) {
697 let j = q * NH + e;
698 let c = ro[I_HCC + j];
699 let k = P(P_HC + j);
700 m += k * rw[S_PHI + c];
701 if (S1 != 0.0) { r += k * rw[S_RATE + c]; }
702 }
703 rw[S_MEL + e] = m;
704 rw[S_MOM + e] = m + 2.0 * S1 * r;
705 if (S1 != 0.0) { pd = 2.0 * S1 * rw[S_RATE + e] * r; }
706 }
707 put(l, 0u, pd);
708 wsum(l, 1u);
709 if (l == 0u) { rw[Q_HD + wg.x] = wres[0]; }
710}
712// accelerations (before rigid removal); per-workgroup damping and external
713// power and sums for the rigid part of the accelerations
714@compute @workgroup_size(64)
715fn beadForces(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
716 /*SGI*/
717 let i = wg.x * 64u + l;
718 let t = f32(count(44u)) * DT;
719 var pd = 0.0;
720 var pe = 0.0;
721 var rig = rigZero();
722 if (i < N) {
723 let im = P(P_IM + i);
724 let ui = gU(i);
725 var ai = vec3f(0.0);
726 if (im != 0.0) {
727 let vi = gV(i);
728 var fi = vec3f(0.0);
729 for (var q = 0u; q < MAXS; q++) {
730 let j = q * N + i;
731 let o = ro[I_SO + j];
732 let b = P_SF + 8u * j;
733 let rd = pv(b + 2u);
734 let du = otherU(o) - ui;
735 let d = rd + du;
736 let len = length(d);
737 let dl = (2.0 * dot(rd, du) + dot(du, du)) / (len + P(b + 5u));
738 fi += (P(b) * (dl + P(b + 6u)) / len) * d;
739 if (ANY_DASH) {
740 let c = P(b + 1u);
741 let vr = otherV(o) - vi;
742 fi += c * vr;
743 pd += P(b + 7u) * c * dot(vr, vr);
744 }
745 }
746 if (NH > 0u) {
747 for (var q = 0u; q < MAXBH; q++) {
748 let ent = ro[I_BHE + q * N + i];
749 let e = ent >> 2u;
750 fi -= rw[S_MOM + e] * ld(S_G + 12u * e + 3u * (ent & 3u));
751 }
752 }
753 if (HAS_MALLET) { fi -= P(P_PW + i) * s3(9u); }
754 if (S0 != 0.0) {
755 let m2 = 2.0 * S0 * P(P_MASS + i);
756 fi -= m2 * vi;
757 pd += m2 * dot(vi, vi);
758 }
759 if (HAS_PULSE && t >= 0.0 && t <= TAU) {
760 let fe = F0 * 0.5 * (1.0 - cos(6.283185307 * t / TAU)) * pv(P_FS + 3u * i);
761 fi += fe;
762 pe += dot(fe, vi);
763 }
764 ai = fi * im;
765 }
766 sA(i, ai);
767 if (FREE_RIGID) { rig = rigAdd(rig, P(P_MASS + i), ui, pv(P_R0 + 3u * i) + ui, ai); }
768 }
769 rigPut(l, rig);
770 put(l, 15u, pd);
771 put(l, 16u, pe);
772 wsum(l, 17u);
773 if (l < 17u) { rw[Q_BF + wg.x * 17u + l] = wres[l]; }
774}
776// remove the rigid part of the accelerations, second half kick; per-workgroup
777// sums for the rigid part of the velocities. Workgroup 0: the mallet's second
778// half kick, dissipated energy and work in.
779@compute @workgroup_size(64)
780fn kick(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
781 /*SGI*/
782 if (FREE_RIGID || wg.x == 0u) { accelRigid(l); }
783 if (wg.x == 0u && l == 0u) {
784 var hd = 0.0;
785 for (var w = 0u; w < NHW; w++) { hd += rw[Q_HD + w]; }
786 if (HAS_MALLET) { ss3(3u, s3(3u) + 0.5 * DT * s3(6u)); }
787 kahan(17u, DT * (wcor[16] + hd + rw[SC + 14u]));
788 kahan(19u, DT * wcor[17]);
789 }
790 let i = wg.x * 64u + l;
791 var p = vec3f(0.0);
792 var lm = vec3f(0.0);
793 if (i < N) {
794 var a = gA(i);
795 var v = gV(i);
796 if (FREE_RIGID) {
797 let s = pos(i);
798 a -= rigidAt(s);
799 sA(i, a);
800 if (P(P_IM + i) != 0.0) { v += 0.5 * DT * a; }
801 let m = P(P_MASS + i);
802 p = m * v;
803 lm = m * cross(s, v);
804 } else if (P(P_IM + i) != 0.0) {
805 v += 0.5 * DT * a;
806 }
807 sV(i, v);
808 }
809 if (FREE_RIGID) {
810 put3(l, 0u, p);
811 put3(l, 3u, lm);
812 wsum(l, 6u);
813 if (l < 6u) { rw[Q_KC + wg.x * 6u + l] = wres[l]; }
814 }
815}
817// per-workgroup audio and (at frames) energy and largest strain for the next
818// sample: bead workgroups first, then hinge workgroups (hinge energy). With
819// correct = true the accelerations still hold their rigid part (forces were
820// just recomputed) and it is removed first.
821fn record(wg: u32, l: u32, correct: bool) {
822 let smp = count(45u);
823 let isFrame = smp % SPF == 0u && smp / SPF < NF;
824 var ap = 0.0;
825 var ep = 0.0;
826 var mx = 0.0;
827 if (wg < NBW) {
828 if (correct) { accelRigid(l); }
829 let i = wg * 64u + l;
830 if (i < N) {
831 var a = gA(i);
832 if (correct) {
833 a -= rigidAt(pos(i));
834 sA(i, a);
835 }
836 let hb = S_HIST + i * H;
837 rw[hb + smp % H] = dot(pv(P_W + 3u * i), a);
838 let i0 = ro[I_I0 + i];
839 let fr = P(P_FR + i);
840 if (smp >= i0) { ap += (1.0 - fr) * rw[hb + (smp - i0) % H]; }
841 if (smp >= i0 + 1u) { ap += fr * rw[hb + (smp - i0 - 1u) % H]; }
842 if (isFrame) {
843 let ui = gU(i);
844 let fo = OFR + (smp / SPF) * 3u * N + 3u * i;
845 ob[fo] = ui.x;
846 ob[fo + 1u] = ui.y;
847 ob[fo + 2u] = ui.z;
848 let vi = gV(i);
849 ep += 0.5 * P(P_MASS + i) * dot(vi, vi);
850 if (P(P_IM + i) != 0.0) {
851 for (var q = 0u; q < MAXS; q++) {
852 let j = q * N + i;
853 let o = ro[I_SO + j];
854 let b = P_SF + 8u * j;
855 let rd = pv(b + 2u);
856 let lr = P(b + 5u);
857 let e0 = P(b + 6u);
858 let du = otherU(o) - ui;
859 let len = length(rd + du);
860 let dl = (2.0 * dot(rd, du) + dot(du, du)) / (len + lr);
861 // energy above the pre-tensioned equilibrium: k/2 ((dl + e0)^2 - e0^2)
862 ep += P(b + 7u) * 0.5 * P(b) * dl * (dl + 2.0 * e0);
863 if (o < N) { mx = max(mx, abs(dl + e0) / (lr - e0)); }
864 }
865 }
866 }
867 }
868 } else if (isFrame) {
869 let e = (wg - NBW) * 64u + l;
870 if (e < NH) { ep = 0.5 * rw[S_PHI + e] * rw[S_MEL + e]; }
871 }
872 put(l, 0u, mx);
873 let largest = wmax(l);
874 put(l, 0u, ap);
875 put(l, 1u, ep);
876 wsum(l, 2u);
877 if (l == 0u) {
878 rw[Q_REC + wg * 3u] = wres[0];
879 rw[Q_REC + wg * 3u + 1u] = wres[1];
880 rw[Q_REC + wg * 3u + 2u] = largest;
881 }
882}
883@compute @workgroup_size(64)
884fn recordSample(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
885 /*SGI*/ record(wg.x, l, false); }
886@compute @workgroup_size(64)
887fn recordSampleCorrect(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
888 /*SGI*/ record(wg.x, l, FREE_RIGID); }
889// writes the last sample
890@compute @workgroup_size(64)
891fn recordLast(@builtin(local_invocation_index) l: u32 /*SG*/) {
892 /*SGI*/ recordSums(l); }
894// at the start: per-workgroup sums for the rigid part of the initial velocities
895@compute @workgroup_size(64)
896fn initRigid(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
897 /*SGI*/
898 let i = wg.x * 64u + l;
899 var rig = rigZero();
900 if (i < N) {
901 let u = gU(i);
902 rig = rigAdd(rig, P(P_MASS + i), u, pv(P_R0 + 3u * i) + u, gV(i));
903 }
904 rigPut(l, rig);
905 wsum(l, 15u);
906 if (l < 15u) { rw[Q_IRV + wg.x * 15u + l] = wres[l]; }
907}
908// and its removal
909@compute @workgroup_size(64)
910fn initVel(@builtin(workgroup_id) wg: vec3u, @builtin(local_invocation_index) l: u32 /*SG*/) {
911 /*SGI*/
912 psum(l, Q_IRV, 15u, 15u, NBW);
913 if (l == 0u) {
914 rigidGeometry();
915 rigidPart(9u);
916 }
917 workgroupBarrier();
918 let i = wg.x * 64u + l;
919 if (i >= N) { return; }
920 sV(i, gV(i) - rigidAt(pos(i)));
921}
922`;
924const ENTRY = ['kickDriftRec', 'kickDrift', 'velCorrect', 'hinges', 'moments', 'beadForces', 'kick', 'recordSample', 'recordSampleCorrect', 'recordLast', 'initRigid', 'initVel'] as const;
925type Entry = (typeof ENTRY)[number];
927export async function simulateGPU(
928 sys: BeadSystem,
929 exc: Excitation,
930 opt: SimOptions,
931 progress?: Progress,
932): Promise<SimResult & { device: string; loopSeconds: number }> {
933 const t0 = performance.now();
934 const gpu = navigator.gpu;
935 if (!gpu) throw new Error('WebGPU is not available in this browser');
936 const adapter = await gpu.requestAdapter({ powerPreference: 'high-performance' });
937 if (!adapter) throw new Error('no WebGPU adapter');
938 const lim = adapter.limits;
939 const sg = adapter.features.has('subgroups');
940 const pk = pack(sys, exc, opt, sg);
941 const biggest = Math.max(pk.ro.byteLength, pk.rw.byteLength, pk.outLen * 4);
942 if (biggest > lim.maxStorageBufferBindingSize || biggest > lim.maxBufferSize)
943 throw new Error(`a GPU buffer would need ${(biggest / 1e6).toFixed(0)} MB (limit ${(Math.min(lim.maxStorageBufferBindingSize, lim.maxBufferSize) / 1e6).toFixed(0)} MB); lower the replay frame rate or the duration`);
944 const device = await adapter.requestDevice({
945 requiredFeatures: sg ? ['subgroups' as GPUFeatureName] : [],
946 requiredLimits: {
947 maxStorageBufferBindingSize: Math.max(biggest, 1 << 20),
948 maxBufferSize: Math.max(biggest, 1 << 20),
949 },
950 });
951 let lost = '';
952 device.lost.then((info) => (lost = info.message || 'GPU device lost'));
953 try {
954 device.pushErrorScope('validation');
955 const module = device.createShaderModule({ code: pk.wgsl });
956 const info = await module.getCompilationInfo();
957 const errs = info.messages.filter((m) => m.type === 'error');
958 if (errs.length) throw new Error(`shader: ${errs.map((m) => `${m.lineNum}:${m.linePos} ${m.message}`).join('; ')}`);
959 const bgl = device.createBindGroupLayout({
960 entries: (['read-only-storage', 'storage', 'storage'] as const).map((type, binding) => ({ binding, visibility: GPUShaderStage.COMPUTE, buffer: { type } })),
961 });
962 const layout = device.createPipelineLayout({ bindGroupLayouts: [bgl] });
963 const pipes = {} as Record<Entry, GPUComputePipeline>;
964 await Promise.all(ENTRY.map(async (e) => (pipes[e] = await device.createComputePipelineAsync({ layout, compute: { module, entryPoint: e } }))));
965 const mk = (data: Float32Array | Uint32Array | null, size: number) => {
966 const buf = device.createBuffer({ size: Math.max(16, size), usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC });
967 if (data) device.queue.writeBuffer(buf, 0, data as Float32Array<ArrayBuffer>);
968 return buf;
969 };
970 const outB = mk(null, pk.outLen * 4);
971 const bind = device.createBindGroup({
972 layout: bgl,
973 entries: [mk(pk.ro, pk.ro.byteLength), mk(pk.rw, pk.rw.byteLength), outB].map((buffer, binding) => ({ binding, resource: { buffer } })),
974 });
976 const { nBead, nHinge, hasMallet, freeRigid, velForces, stepsPerSample } = pk;
977 const nh = sys.bending ? sys.bending.w0.length : 0;
978 type D = [Entry, number];
979 const forces: D[] = [];
980 if (nHinge + (hasMallet ? 1 : 0) > 0) forces.push(['hinges', nHinge + (hasMallet ? 1 : 0)]);
981 if (nh > 0) forces.push(['moments', nHinge]);
982 forces.push(['beadForces', nBead]);
983 // one sample: record it, then the steps to the next one; the accelerations
984 // hold their rigid part after a force evaluation until a kick removes it
985 const sample = (first: boolean, last: boolean): D[] => {
986 const raw = first || velForces;
987 const list: D[] = [[raw ? 'recordSampleCorrect' : 'recordSample', nBead + nHinge]];
988 if (last) return [...list, ['recordLast', 1]];
989 for (let j = 0; j < stepsPerSample; j++) list.push([j === 0 ? 'kickDriftRec' : 'kickDrift', nBead], ...forces, ['kick', nBead]);
990 if (freeRigid) list.push(['velCorrect', nBead]);
991 if (velForces) list.push(...forces);
992 return list;
993 };
994 const start: D[] = [...(freeRigid ? ([['initRigid', nBead], ['initVel', nBead]] as D[]) : []), ...forces];
995 const encode = (pass: GPUComputePassEncoder, list: D[]) => {
996 for (const [e, w] of list) {
997 pass.setPipeline(pipes[e]);
998 pass.dispatchWorkgroups(w);
999 }
1000 };
1001 const submit = (fill: (pass: GPUComputePassEncoder) => void) => {
1002 const enc = device.createCommandEncoder();
1003 const pass = enc.beginComputePass();
1004 pass.setBindGroup(0, bind);
1005 fill(pass);
1006 pass.end();
1007 device.queue.submit([enc.finish()]);
1008 return device.queue.onSubmittedWorkDone();
1009 };
1010 await submit((pass) => encode(pass, start));
1011 const err = await device.popErrorScope();
1012 if (err) throw new Error(err.message);
1014 // chunks of samples, about 50 ms of GPU work each, two in flight so the
1015 // GPU does not wait while the next one is encoded
1016 const tLoop = performance.now();
1017 const mid = sample(false, false);
1018 let chunk = 16;
1019 let pending: Promise<void> | null = null;
1020 let pendingEnd = 0;
1021 let tPrev = performance.now();
1022 for (let s = 0; s < pk.nSamples; ) {
1023 const e = Math.min(pk.nSamples, s + chunk);
1024 const done = submit((pass) => {
1025 for (let k = s; k < e; k++) encode(pass, k === 0 || k === pk.nSamples - 1 ? sample(k === 0, k === pk.nSamples - 1) : mid);
1026 });
1027 if (pending) {
1028 await pending;
1029 if (lost) throw new Error(lost);
1030 const now = performance.now();
1031 chunk = Math.max(1, Math.min(Math.round(chunk * Math.min(2, 50 / Math.max(now - tPrev, 1))), 4096));
1032 tPrev = now;
1033 progress?.(pendingEnd / pk.nSamples);
1034 }
1035 pending = done;
1036 pendingEnd = e;
1037 s = e;
1038 }
1039 await pending;
1040 if (lost) throw new Error(lost);
1041 progress?.(1);
1042 const loopSeconds = (performance.now() - tLoop) / 1000;
1044 const stage = device.createBuffer({ size: Math.max(16, pk.outLen * 4), usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST });
1045 const enc = device.createCommandEncoder();
1046 enc.copyBufferToBuffer(outB, 0, stage, 0, Math.max(16, pk.outLen * 4));
1047 device.queue.submit([enc.finish()]);
1048 await stage.mapAsync(GPUMapMode.READ);
1049 const out = new Float32Array(stage.getMappedRange().slice(0, pk.outLen * 4));
1050 stage.destroy();
1051 if (lost) throw new Error(lost);
1053 const { nSamples, nFrames, samplesPerFrame, dt } = pk;
1054 const audio = out.slice(0, nSamples);
1055 for (let k = 0; k < audio.length; k++) if (!Number.isFinite(audio[k])) throw new Error('the GPU run blew up (non-finite output)');
1056 const contactForce = exc.mallet ? out.slice(nSamples, 2 * nSamples) : undefined;
1057 const O_FSC = 2 * nSamples;
1058 const O_MP = O_FSC + 6 * nFrames;
1059 const frameTimes = new Float64Array(nFrames);
1060 const energy = new Float64Array(nFrames);
1061 const balance = new Float64Array(nFrames);
1062 let maxStretch = 0;
1063 let diss = 0;
1064 let work = 0;
1065 let E0 = 0;
1066 for (let k = 0; k < nFrames; k++) {
1067 const b = O_FSC + 6 * k;
1068 frameTimes[k] = k * samplesPerFrame * stepsPerSample * dt;
1069 energy[k] = out[b];
1070 const total = out[b] + out[b + 1];
1071 diss += out[b + 2];
1072 work += out[b + 3];
1073 if (k === 0) E0 = total;
1074 balance[k] = total + diss - work - E0;
1075 maxStretch = Math.max(maxStretch, out[b + 4]);
1076 }
1077 return {
1078 audio,
1079 fs: opt.fs,
1080 frames: out.slice(pk.framesAt),
1081 frameTimes,
1082 energy,
1083 balance,
1084 dt,
1085 stepsPerSample,
1086 omegaMax: pk.omegaMax,
1087 wallSeconds: (performance.now() - t0) / 1000,
1088 maxStretch,
1089 contactForce,
1090 malletPos: exc.mallet ? out.slice(O_MP, O_MP + 3 * nFrames) : undefined,
1091 loopSeconds,
1092 device: [adapter.info?.vendor, adapter.info?.architecture].filter(Boolean).join(' ') || 'GPU',
1093 };
1094 } finally {
1095 device.destroy();
1096 }
1097}