917e5cdWebGPU toggle: single-precision GPU integration (displacement-based, multi-dispatch phases, frames read back), linear modes behind a toggle (off by default), shared prepare()Jeremy Magland 1// Bead solver on the GPU: the same scheme as solver.ts (velocity Verlet or
2// its Yoshida composition, with exact damping applied once per audio sample),
3// in single precision.
4//
5// Each phase of a step is its own dispatch over beads or elements; the host
6// encodes the phases for a chunk of audio samples into one command buffer.
7// Per-phase constants (kick and drift lengths, Taylor coefficients) come from
8// a table selected with a dynamic uniform offset, and the current sample
9// index lives in a GPU buffer, which the record phase advances.
10//
11// To keep single precision accurate, the state is the displacement u from
12// the rest positions, never the absolute position, and every element works
13// from its rest geometry (computed in double precision on the CPU) plus
14// displacement differences. Element forces are written to per-element slots
15// and gathered per bead through an incidence list, since there are no atomic
16// float additions.
18const WG: u32 = 64u;
20struct Hdr {
21 n: u32, ns: u32, nb: u32, nh: u32,
22 nd: u32, nl: u32, quantity: u32, npulse: u32,
23 sample0: u32, frameStride: u32, frameSize: u32, nImp: u32,
24 impOff: u32, oFrames: u32, pad0: u32, pad1: u32,
25 // offsets into P (f32 parameters)
26 oInvm: u32, oMass: u32, oFric: u32, oS: u32,
27 oB: u32, oH: u32, oD: u32, oPulse: u32,
28 // offsets into I (u32 indices)
29 iS: u32, iB: u32, iH: u32, iD: u32,
30 iFr: u32, iFe: u32, iDr: u32, iDe: u32,
31 iPr: u32, iPe: u32, iL: u32, iImp: u32,
32 // offsets into E for the bending and hinge slots (forces / damping)
33 eFb: u32, eFh: u32, eDb: u32, eDh: u32,
34 Delta: f32, dirx: f32, diry: f32, dirz: f32,
35}
37// per-phase constants
38struct Ph {
39 kick: f32, drift: f32, tOff: f32, gather: u32,
40 c: f32, saveAv: u32, last: u32, pad: u32,
41}
43@group(0) @binding(0) var<uniform> H: Hdr;
44@group(0) @binding(1) var<storage, read> P: array<f32>;
45@group(0) @binding(2) var<storage, read> I: array<u32>;
46@group(0) @binding(3) var<storage, read_write> S: array<f32>; // u | v | F | T | Av | KE0 | diss | theta0
47@group(0) @binding(4) var<storage, read_write> E: array<f32>; // element slots
48@group(0) @binding(5) var<storage, read_write> O: array<f32>; // audio | frames
49@group(0) @binding(6) var<storage, read_write> C: array<u32>; // [current sample]
50@group(1) @binding(0) var<uniform> Q: Ph;
52fn offV() -> u32 { return 3u * H.n; }
53fn offF() -> u32 { return 6u * H.n; }
54fn offT() -> u32 { return 9u * H.n; }
55fn offAv() -> u32 { return 12u * H.n; }
56fn offKE() -> u32 { return 15u * H.n; }
57fn offDiss() -> u32 { return 16u * H.n; }
58fn offTh0() -> u32 { return 17u * H.n; }
60fn S3(o: u32) -> vec3<f32> { return vec3<f32>(S[o], S[o + 1u], S[o + 2u]); }
61fn setS3(o: u32, w: vec3<f32>) { S[o] = w.x; S[o + 1u] = w.y; S[o + 2u] = w.z; }
62fn setE3(o: u32, w: vec3<f32>) { E[o] = w.x; E[o + 1u] = w.y; E[o + 2u] = w.z; }
63fn E3(o: u32) -> vec3<f32> { return vec3<f32>(E[o], E[o + 1u], E[o + 2u]); }
64fn P3(o: u32) -> vec3<f32> { return vec3<f32>(P[o], P[o + 1u], P[o + 2u]); }
65fn U(b: u32) -> vec3<f32> { return S3(3u * b); }
66fn V(b: u32) -> vec3<f32> { return S3(offV() + 3u * b); }
68// Hinge h: angle and gradient from the rest vectors relative to bead j plus
69// displacement differences.
70struct HingeG { theta: f32, g: array<vec3<f32>, 4> }
72fn hingeAt(h: u32) -> HingeG {
73 let bi = I[H.iH + 4u * h];
74 let bj = I[H.iH + 4u * h + 1u];
75 let bk = I[H.iH + 4u * h + 2u];
76 let bl = I[H.iH + 4u * h + 3u];
77 let o = H.oH + 16u * h;
78 let uj = U(bj);
79 let pa = P3(o) + (U(bi) - uj); // bead i
80 let pk = P3(o + 3u) + (U(bk) - uj); // bead k
81 let pl = P3(o + 6u) + (U(bl) - uj); // bead l; bead j is the origin
82 let e = pk;
83 let eLen = length(e);
84 let b = pa - pk;
85 let c = pl - pk;
86 let n1 = cross(pa, b);
87 let n2 = cross(c, pl);
88 let n1sq = dot(n1, n1);
89 let n2sq = dot(n2, n2);
90 var r: HingeG;
91 r.theta = -atan2(dot(cross(n1, n2), e) / eLen, dot(n1, n2));
92 let p1 = dot(pa, e) / eLen;
93 let q1 = dot(b, e) / eLen;
94 let p2 = dot(pl, e) / eLen;
95 let q2 = dot(c, e) / eLen;
96 let m1 = 1.0 / n1sq;
97 let m2 = 1.0 / n2sq;
98 r.g[0] = (eLen * m1) * n1;
99 r.g[1] = q1 * m1 * n1 + q2 * m2 * n2;
100 r.g[2] = -p1 * m1 * n1 - p2 * m2 * n2;
101 r.g[3] = (eLen * m2) * n2;
102 return r;
103}
105fn wrapAngle(d: f32) -> f32 {
106 var x = d;
107 if (x > 3.14159265) { x -= 6.28318531; }
108 if (x <= -3.14159265) { x += 6.28318531; }
109 return x;
110}
112// ---- forces ---------------------------------------------------------------------
114// one invocation per element: springs, then bending triplets, then hinges
115@compute @workgroup_size(WG)
116fn elements(@builtin(global_invocation_id) gid: vec3<u32>) {
117 var q = gid.x;
118 if (q < H.ns) {
119 let o = H.oS + 8u * q;
120 let r0 = P3(o);
121 let du = U(I[H.iS + 2u * q + 1u]) - U(I[H.iS + 2u * q]);
122 let r = r0 + du;
123 let L = length(r);
124 // L - |r0| without cancellation, plus the rest stretch |r0| - L0
125 let e = P[o + 4u] + (2.0 * dot(r0, du) + dot(du, du)) / (L + P[o + 3u]);
126 let f = P[o + 5u] * e + P[o + 6u] * e * e * e;
127 setE3(3u * q, (f / L) * r);
128 return;
129 }
130 q -= H.ns;
131 if (q < H.nb) {
132 let o = H.oB + 8u * q;
133 let D = P3(o) + U(I[H.iB + 3u * q]) - 2.0 * U(I[H.iB + 3u * q + 1u]) + U(I[H.iB + 3u * q + 2u]);
134 setE3(H.eFb + 3u * q, P[o + 3u] * D);
135 return;
136 }
137 q -= H.nb;
138 if (q < H.nh) {
139 let hg = hingeAt(q);
140 let s = P[H.oH + 16u * q + 9u] * wrapAngle(hg.theta - S[offTh0() + q]);
141 for (var p = 0u; p < 4u; p++) { setE3(H.eFh + 12u * q + 3u * p, s * hg.g[p]); }
142 }
143}
145fn pulseForce(b: u32, t: f32) -> vec3<f32> {
146 var f = vec3<f32>(0.0);
147 for (var q = I[H.iPr + b]; q < I[H.iPr + b + 1u]; q++) {
148 let o = H.oPulse + 8u * I[H.iPe + q];
149 let x = (t - P[o]) / P[o + 1u];
150 if (x > 0.0 && x < 1.0) { f += (0.5 * (1.0 - cos(6.28318531 * x))) * P3(o + 2u); }
151 }
152 return f;
153}
155// gather (optional), kick and drift; time t = (sample index) * Delta + tOff
156@compute @workgroup_size(WG)
157fn beads(@builtin(global_invocation_id) gid: vec3<u32>) {
158 let b = gid.x;
159 if (b >= H.n) { return; }
160 var F = S3(offF() + 3u * b);
161 if (Q.gather == 1u) {
162 F = vec3<f32>(0.0);
163 for (var q = I[H.iFr + b]; q < I[H.iFr + b + 1u]; q++) {
164 F += bitcast<f32>(I[H.iFe + 2u * q + 1u]) * E3(I[H.iFe + 2u * q]);
165 }
166 if (H.npulse > 0u) { F += pulseForce(b, f32(C[0]) * H.Delta + Q.tOff); }
167 setS3(offF() + 3u * b, F);
168 }
169 let v = V(b) + (Q.kick * P[H.oInvm + b]) * F;
170 setS3(offV() + 3u * b, v);
171 setS3(3u * b, U(b) + Q.drift * v);
172}
174// ---- sample boundary ----------------------------------------------------------------
176// frame copy (if this sample is a frame), T = v, kinetic energy before damping
177@compute @workgroup_size(WG)
178fn pre(@builtin(global_invocation_id) gid: vec3<u32>) {
179 let b = gid.x;
180 if (b >= H.n) { return; }
181 let sample = C[0];
182 let v = V(b);
183 if (sample % H.frameStride == 0u) {
184 let first = ((H.sample0 + H.frameStride - 1u) / H.frameStride) * H.frameStride;
185 let o = H.oFrames + ((sample - first) / H.frameStride) * H.frameSize;
186 let u = U(b);
187 O[o + 3u * b] = u.x;
188 O[o + 3u * b + 1u] = u.y;
189 O[o + 3u * b + 2u] = u.z;
190 O[o + 3u * H.n + 3u * b] = v.x;
191 O[o + 3u * H.n + 3u * b + 1u] = v.y;
192 O[o + 3u * H.n + 3u * b + 2u] = v.z;
193 O[o + 6u * H.n + b] = S[offDiss() + b];
194 }
195 setS3(offT() + 3u * b, v);
196 S[offKE() + b] = 0.5 * P[H.oMass + b] * dot(v, v);
197}
199// T = v, for each further piece of a damping step split into pieces
200@compute @workgroup_size(WG)
201fn resetT(@builtin(global_invocation_id) gid: vec3<u32>) {
202 if (gid.x < H.n) { setS3(offT() + 3u * gid.x, V(gid.x)); }
203}
205// one Taylor term of exp(-tau A): the damping forces of the elements from T
206@compute @workgroup_size(WG)
207fn dampElements(@builtin(global_invocation_id) gid: vec3<u32>) {
208 var q = gid.x;
209 if (q < H.nd) {
210 let bi = I[H.iD + 2u * q];
211 let bj = I[H.iD + 2u * q + 1u];
212 let o = H.oD + 8u * q;
213 let rw = S3(offT() + 3u * bi) - S3(offT() + 3u * bj);
214 var f = P[o + 3u] * rw;
215 if (P[o + 4u] != 0.0) {
216 let dir = normalize(P3(o) + U(bj) - U(bi));
217 f += (P[o + 4u] * dot(dir, rw)) * dir;
218 }
219 setE3(3u * q, f);
220 return;
221 }
222 q -= H.nd;
223 if (q < H.nb) {
224 let cb = P[H.oB + 8u * q + 4u];
225 let w = S3(offT() + 3u * I[H.iB + 3u * q]) - 2.0 * S3(offT() + 3u * I[H.iB + 3u * q + 1u]) + S3(offT() + 3u * I[H.iB + 3u * q + 2u]);
226 setE3(H.eDb + 3u * q, cb * w);
227 return;
228 }
229 q -= H.nb;
230 if (q < H.nh) {
231 let ch = P[H.oH + 16u * q + 10u];
232 if (ch == 0.0) { return; } // no damping entries reference these slots
233 let hg = hingeAt(q);
234 var s = 0.0;
235 for (var p = 0u; p < 4u; p++) { s += dot(hg.g[p], S3(offT() + 3u * I[H.iH + 4u * q + p])); }
236 for (var p = 0u; p < 4u; p++) { setE3(H.eDh + 12u * q + 3u * p, (ch * s) * hg.g[p]); }
237 }
238}
240// A T per bead, then T <- c A T and v <- v + T (c = -tau/k)
241@compute @workgroup_size(WG)
242fn dampBeads(@builtin(global_invocation_id) gid: vec3<u32>) {
243 let b = gid.x;
244 if (b >= H.n) { return; }
245 var f = P[H.oFric + b] * S3(offT() + 3u * b);
246 for (var q = I[H.iDr + b]; q < I[H.iDr + b + 1u]; q++) {
247 f += bitcast<f32>(I[H.iDe + 2u * q + 1u]) * E3(I[H.iDe + 2u * q]);
248 }
249 let a = P[H.oInvm + b] * f;
250 if (Q.saveAv == 1u) { setS3(offAv() + 3u * b, a); }
251 let t = Q.c * a;
252 setS3(offT() + 3u * b, t);
253 let v = V(b) + t;
254 setS3(offV() + 3u * b, v);
255 if (Q.last == 1u) {
256 S[offDiss() + b] += S[offKE() + b] - 0.5 * P[H.oMass + b] * dot(v, v);
257 }
258}
260// the listener's sample; advances the sample index
261@compute @workgroup_size(1)
262fn record() {
263 let dir = vec3<f32>(H.dirx, H.diry, H.dirz);
264 var val = 0.0;
265 for (var q = 0u; q < H.nl; q++) {
266 let b = I[H.iL + q];
267 if (H.quantity == 0u) { val += dot(U(b), dir); }
268 else if (H.quantity == 1u) { val += dot(S3(offT() + 3u * b), dir); } // T holds v at the boundary
269 else { val += dot(P[H.oInvm + b] * S3(offF() + 3u * b) - S3(offAv() + 3u * b), dir); }
270 }
271 O[C[0] - H.sample0] = val / f32(H.nl);
272 C[0] = C[0] + 1u;
273}
275// velocity impulses at the start of a chunk: [bead, vx, vy, vz] per entry
276@compute @workgroup_size(WG)
277fn impulse(@builtin(global_invocation_id) gid: vec3<u32>) {
278 let q = gid.x;
279 if (q >= H.nImp) { return; }
280 let o = H.iImp + 4u * (H.impOff + q);
281 let b = I[o];
282 if (P[H.oInvm + b] > 0.0) {
283 setS3(offV() + 3u * b, V(b) + vec3<f32>(bitcast<f32>(I[o + 1u]), bitcast<f32>(I[o + 2u]), bitcast<f32>(I[o + 3u])));
284 }
285}
287// rest angles of the hinges (u = 0)
288@compute @workgroup_size(WG)
289fn initHinges(@builtin(global_invocation_id) gid: vec3<u32>) {
290 if (gid.x < H.nh) { S[offTh0() + gid.x] = hingeAt(gid.x).theta; }
291}