1// WGSL kernels of the time loop. Four entry points: the element pass (forces, sigma1 term,
2// contact forces of the hits), the node pass (gather, contact, lagged rigid-mode projection,
3// leapfrog update, partial sums, exciter update, and the normal acceleration with its decimation
4// accumulators), the gather pass (element means of a finished audio sample into the history,
5// then the retarded-time sum at the listener, one partial sum per workgroup) and the capture pass
6// (displacement frames). Sizes
7// and constants are baked into the source. See THEORY.md for the buffer layouts and the ordering
8// argument behind the two step counters.
10export interface ShaderConstants {
11 N: number; E: number; K: number; P: number; M: number; W: number; D: number; G: number; T: number;
12 accRows: number; dt: number; sigma0: number; nl: number;
13}
15export const ELEM_STRIDE = 96; // Gx[3] Gy[3] A z0x z0y C11 C12 C33 cd Ke[81] pad[2]
16export const NODE_STRIDE = 44; // Minv[5] M[5] Rm[6][5] normal[3] pad[1]; the decimation taps follow the nodes
17export const HIT_HEADER = 12; // kPoint alpha mass gap velocity zs vs F Fmax pad[3]
18export const HITU_HEADER = 4; // releaseStep ncont pad[2]
19export const SYNC_HEADER = 4; // stepA stepB pad[2], then partials[2][G][12]
20export const RADCOEF_STRIDE = 36; // coefF coefB i0F i0B tapsF[16] tapsB[16]
21export const WORKGROUP_NODES = 256;
22export const GATHER_ELEMS = 256; // elements per gather workgroup
23export const STATE_FLOATS = 13; // per node: q[5] v[5] vPrev[3]
24/** Chunk uniform: base size frameBase every accBase accRows pad pad (eight 32-bit words). */
25export const CHUNK_WORDS = 8;
27function f(x: number): string {
28 const s = x.toPrecision(9);
29 return /[.e]/.test(s) ? s : s + ".0";
30}
32export function shaderSource(c: ShaderConstants): string {
33 const HF = HIT_HEADER + 5 * c.P, HU = HITU_HEADER + 3 * c.P;
34 return /* wgsl */ `
35const N: u32 = ${c.N}u; const E: u32 = ${c.E}u; const K: u32 = ${c.K}u; const P: u32 = ${c.P}u;
36const M: u32 = ${c.M}u; const W: u32 = ${c.W}u; const D: u32 = ${c.D}u; const G: u32 = ${c.G}u; const T: u32 = ${c.T}u;
37const DT: f32 = ${f(c.dt)}; const DAMP0: f32 = ${f(1 / (1 + 2 * c.sigma0 * c.dt))}; const NL: f32 = ${f(c.nl)};
38const EF: u32 = ${ELEM_STRIDE}u; const NF: u32 = ${NODE_STRIDE}u; const HF: u32 = ${HF}u; const HU: u32 = ${HU}u; const RC: u32 = ${RADCOEF_STRIDE}u;
39const OFF_V: u32 = 5u * N; const OFF_VWP: u32 = 10u * N;
40const OFF_NHIST: u32 = W * N; const OFF_AHIST: u32 = (W + D) * N; const OFF_ACCOUT: u32 = (W + D + M) * N;
41const GE: u32 = ${Math.ceil(c.E / GATHER_ELEMS)}u;
42const OFF_TAPS: u32 = NF * N; const OFF_PART: u32 = ${SYNC_HEADER}u;
44struct Chunk { base: i32, size: u32, frameBase: u32, every: u32, accBase: i32, accRows: u32, pad0: u32, pad1: u32 }
46// ------------------------------------------------------------------ element pass
47@group(0) @binding(0) var<storage, read> elemF: array<f32>;
48@group(0) @binding(1) var<storage, read> elemU: array<u32>;
49@group(0) @binding(2) var<storage, read> stateR: array<f32>;
50@group(0) @binding(3) var<storage, read_write> elemForce: array<f32>;
51@group(0) @binding(4) var<storage, read_write> hitsF: array<f32>;
52@group(0) @binding(5) var<storage, read> hitsUR: array<u32>;
53@group(0) @binding(6) var<storage, read_write> syncE: array<f32>;
55@compute @workgroup_size(64)
56fn elementPass(@builtin(global_invocation_id) gid: vec3<u32>) {
57 let id = gid.x;
58 let step = bitcast<u32>(syncE[0]);
59 if (id == 0u) { syncE[1] = syncE[0]; }
60 if (id < E) {
61 let e = id;
62 let b = e * EF;
63 let n0 = elemU[3u * e]; let n1 = elemU[3u * e + 1u]; let n2 = elemU[3u * e + 2u];
64 var qe: array<f32, 15>;
65 for (var d = 0u; d < 5u; d++) { qe[d] = stateR[5u * n0 + d]; qe[5u + d] = stateR[5u * n1 + d]; qe[10u + d] = stateR[5u * n2 + d]; }
66 let gx = vec3<f32>(elemF[b], elemF[b + 1u], elemF[b + 2u]);
67 let gy = vec3<f32>(elemF[b + 3u], elemF[b + 4u], elemF[b + 5u]);
68 let A = elemF[b + 6u]; let z0x = elemF[b + 7u]; let z0y = elemF[b + 8u];
69 let C11 = elemF[b + 9u]; let C12 = elemF[b + 10u]; let C33 = elemF[b + 11u]; let cd = elemF[b + 12u];
70 let uu = vec3<f32>(qe[0], qe[5], qe[10]); let vv = vec3<f32>(qe[1], qe[6], qe[11]); let ww = vec3<f32>(qe[2], qe[7], qe[12]);
71 let ux = dot(gx, uu); let uy = dot(gy, uu); let vx = dot(gx, vv); let vy = dot(gy, vv); let wx = dot(gx, ww); let wy = dot(gy, ww);
72 let ax = NL * wx + z0x; let ay = NL * wy + z0y;
73 let exx = ux + NL * 0.5 * wx * wx + z0x * wx;
74 let eyy = vy + NL * 0.5 * wy * wy + z0y * wy;
75 let gxy = uy + vx + NL * wx * wy + z0x * wy + z0y * wx;
76 let Nxx = C11 * exx + C12 * eyy; let Nyy = C12 * exx + C11 * eyy; let Nxy = C33 * gxy;
77 let vw = vec3<f32>(stateR[OFF_V + 5u * n0 + 2u], stateR[OFF_V + 5u * n1 + 2u], stateR[OFF_V + 5u * n2 + 2u]);
78 let vwx = dot(gx, vw); let vwy = dot(gy, vw);
79 var out: array<f32, 15>;
80 for (var i = 0u; i < 3u; i++) {
81 let g_x = gx[i]; let g_y = gy[i];
82 out[5u * i] = -A * (Nxx * g_x + Nxy * g_y);
83 out[5u * i + 1u] = -A * (Nyy * g_y + Nxy * g_x);
84 out[5u * i + 2u] = -A * (Nxx * ax * g_x + Nyy * ay * g_y + Nxy * (ay * g_x + ax * g_y)) - cd * (g_x * vwx + g_y * vwy);
85 out[5u * i + 3u] = 0.0; out[5u * i + 4u] = 0.0;
86 }
87 var U: array<f32, 9>;
88 for (var i = 0u; i < 3u; i++) { U[3u * i] = qe[5u * i + 2u]; U[3u * i + 1u] = qe[5u * i + 3u]; U[3u * i + 2u] = qe[5u * i + 4u]; }
89 for (var i = 0u; i < 9u; i++) {
90 var s = 0.0;
91 for (var j = 0u; j < 9u; j++) { s += elemF[b + 13u + 9u * i + j] * U[j]; }
92 let slot = 5u * (i / 3u) + 2u + (i % 3u);
93 out[slot] -= s;
94 }
95 for (var i = 0u; i < 15u; i++) { elemForce[15u * e + i] = out[i]; }
96 } else if (id < E + K) {
97 // contact forces of hit k at this step
98 let k = id - E; let hf = k * HF; let hu = k * HU;
99 var wp: array<f32, ${c.P}>;
100 var wmax = -1.0e30;
101 for (var p = 0u; p < P; p++) {
102 var w = hitsF[hf + ${HIT_HEADER}u + 3u * P + p];
103 for (var i = 0u; i < 3u; i++) { w += hitsF[hf + ${HIT_HEADER}u + 3u * p + i] * stateR[5u * hitsUR[hu + ${HITU_HEADER}u + 3u * p + i] + 2u]; }
104 wp[p] = w; wmax = max(wmax, w);
105 }
106 var zs = hitsF[hf + 5u];
107 if (step == hitsUR[hu]) { zs = wmax + hitsF[hf + 3u]; hitsF[hf + 5u] = zs; hitsF[hf + 6u] = -hitsF[hf + 4u]; }
108 let isActive = step >= hitsUR[hu];
109 let kPoint = hitsF[hf]; let alpha = hitsF[hf + 1u];
110 var F = 0.0;
111 for (var p = 0u; p < P; p++) {
112 let delta = max(wp[p] - zs, 0.0);
113 var fp = 0.0;
114 if (isActive && delta > 0.0) { fp = kPoint * pow(delta, alpha); }
115 hitsF[hf + ${HIT_HEADER}u + 4u * P + p] = fp; F += fp;
116 }
117 hitsF[hf + 7u] = F;
118 }
119}
121// ------------------------------------------------------------------ node pass
122@group(1) @binding(0) var<storage, read_write> state: array<f32>;
123@group(1) @binding(1) var<storage, read> elemForceR: array<f32>;
124@group(1) @binding(2) var<storage, read> nodeF: array<f32>;
125@group(1) @binding(3) var<storage, read> adj: array<u32>;
126@group(1) @binding(4) var<storage, read_write> hitsFN: array<f32>;
127@group(1) @binding(5) var<storage, read_write> hitsU: array<u32>;
128@group(1) @binding(6) var<storage, read_write> sync: array<f32>;
129@group(1) @binding(7) var<storage, read_write> rad: array<f32>;
130@group(1) @binding(8) var<uniform> chunkN: Chunk;
132var<workgroup> red: array<f32, ${12 * WORKGROUP_NODES}>;
134fn modW(x: i32) -> u32 { return u32(((x % i32(W)) + i32(W)) % i32(W)); }
136@compute @workgroup_size(${WORKGROUP_NODES})
137fn nodePass(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>, @builtin(workgroup_id) wid: vec3<u32>) {
138 let nd = gid.x; let t = lid.x;
139 let step = bitcast<u32>(sync[1]);
140 let parity = step & 1u; let prevp = (step + 1u) & 1u;
141 var c: array<f32, 12>;
142 for (var k = 0u; k < 12u; k++) {
143 var s = 0.0;
144 for (var g = 0u; g < G; g++) { s += sync[OFF_PART + (prevp * G + g) * 12u + k]; }
145 c[k] = s;
146 }
147 var part: array<f32, 12>;
148 for (var k = 0u; k < 12u; k++) { part[k] = 0.0; }
149 if (nd < N) {
150 var fo: array<f32, 5>;
151 for (var d = 0u; d < 5u; d++) { fo[d] = 0.0; }
152 let a0 = adj[nd]; let a1 = adj[nd + 1u];
153 for (var a = a0; a < a1; a++) {
154 let ent = adj[N + 1u + a]; let e = ent / 3u; let i = ent - 3u * e;
155 for (var d = 0u; d < 5u; d++) { fo[d] += elemForceR[15u * e + 5u * i + d]; }
156 }
157 for (var k = 0u; k < K; k++) {
158 for (var p = 0u; p < P; p++) {
159 for (var i = 0u; i < 3u; i++) {
160 if (hitsU[k * HU + ${HITU_HEADER}u + 3u * p + i] == nd) {
161 fo[2] -= hitsFN[k * HF + ${HIT_HEADER}u + 4u * P + p] * hitsFN[k * HF + ${HIT_HEADER}u + 3u * p + i];
162 }
163 }
164 }
165 }
166 let nb = nd * NF;
167 var acc = 0.0;
168 for (var d = 0u; d < 5u; d++) {
169 let idx = 5u * nd + d;
170 var vd = state[OFF_V + idx]; var qd = state[idx];
171 for (var k = 0u; k < 6u; k++) { let r = nodeF[nb + 10u + 5u * k + d]; vd -= c[k] * r; qd -= c[6u + k] * r; }
172 let v1 = (vd + DT * fo[d] * nodeF[nb + d]) * DAMP0;
173 // normal acceleration of the previous step from consecutive projected velocities
174 if (d < 3u) { acc += nodeF[nb + 40u + d] * (vd - state[OFF_VWP + 3u * nd + d]) / DT; state[OFF_VWP + 3u * nd + d] = vd; }
175 let q1 = qd + DT * v1;
176 state[idx] = q1; state[OFF_V + idx] = v1;
177 let mm = nodeF[nb + 5u + d];
178 for (var k = 0u; k < 6u; k++) { let r = nodeF[nb + 10u + 5u * k + d] * mm; part[k] += r * v1; part[6u + k] += r * q1; }
179 }
180 // decimation of the acceleration of sub-step s = step - 1 into the W accumulator slots of this node;
181 // slot w belongs to audio sample j0 + w - W/2 + 1, kept at ring position (sample mod W)
182 if (step >= 1u) {
183 let s = step - 1u; let j0 = s / M; let kk = s - j0 * M;
184 if (kk == 0u) {
185 let jE = i32(j0) - i32(W / 2u);
186 let slotE = modW(jE);
187 let val = rad[slotE * N + nd];
188 if (jE >= 0) {
189 rad[OFF_NHIST + (u32(jE) % D) * N + nd] = val;
190 let row = jE - chunkN.accBase;
191 if (chunkN.accRows > 0u && row >= 0 && u32(row) < chunkN.accRows) { rad[OFF_ACCOUT + u32(row) * N + nd] = val; }
192 }
193 rad[slotE * N + nd] = 0.0;
194 }
195 // keep the sub-step's acceleration; once the M values of the sample are in, add them to the W slots at once
196 rad[OFF_AHIST + kk * N + nd] = acc;
197 if (kk == M - 1u) {
198 let r = modW(i32(j0) + 1 - i32(W / 2u));
199 var sl: array<f32, ${c.W}>;
200 for (var w = 0u; w < W; w++) { sl[w] = 0.0; }
201 for (var m = 0u; m < M; m++) {
202 let a = rad[OFF_AHIST + m * N + nd];
203 for (var w = 0u; w < W; w++) { sl[w] += a * nodeF[OFF_TAPS + w * M + m]; }
204 }
205 for (var w = 0u; w < W; w++) { rad[((r + w) % W) * N + nd] += sl[w]; }
206 }
207 }
208 }
209 for (var k = 0u; k < 12u; k++) { red[t * 12u + k] = part[k]; }
210 workgroupBarrier();
211 for (var s = ${WORKGROUP_NODES / 2}u; s > 0u; s = s >> 1u) {
212 if (t < s) { for (var k = 0u; k < 12u; k++) { red[t * 12u + k] += red[(t + s) * 12u + k]; } }
213 workgroupBarrier();
214 }
215 if (t == 0u) { for (var k = 0u; k < 12u; k++) { sync[OFF_PART + (parity * G + wid.x) * 12u + k] = red[k]; } }
216 if (nd == 0u) {
217 for (var k = 0u; k < K; k++) {
218 let hf = k * HF; let hu = k * HU;
219 let F = hitsFN[hf + 7u];
220 let vs = hitsFN[hf + 6u] + DT * F / hitsFN[hf + 2u];
221 hitsFN[hf + 6u] = vs; hitsFN[hf + 5u] = hitsFN[hf + 5u] + DT * vs;
222 hitsFN[hf + 8u] = max(hitsFN[hf + 8u], F);
223 if (F > 0.0) { hitsU[hu + 1u] = hitsU[hu + 1u] + 1u; }
224 }
225 sync[0] = bitcast<f32>(step + 1u);
226 }
227}
229// ------------------------------------------------------------------ gather pass (listener), once per finished audio sample
230@group(2) @binding(0) var<storage, read> syncG: array<f32>;
231@group(2) @binding(1) var<storage, read> radR: array<f32>;
232@group(2) @binding(2) var<storage, read> radCoef: array<f32>;
233@group(2) @binding(3) var<storage, read_write> pressure: array<f32>;
234@group(2) @binding(4) var<uniform> chunk: Chunk;
235@group(2) @binding(5) var<storage, read_write> elemHist: array<f32>;
236@group(2) @binding(6) var<storage, read> elemUG: array<u32>;
238var<workgroup> redF: array<f32, ${GATHER_ELEMS}>;
239var<workgroup> redB: array<f32, ${GATHER_ELEMS}>;
241fn delayed(e: u32, j: i32, i0: i32, tapBase: u32) -> f32 {
242 var s = 0.0;
243 for (var t = 0u; t < T; t++) {
244 let lag = i0 + i32(t) - i32(T / 2u) + 1; let jj = j - lag;
245 if (jj >= 0 && lag >= 0) { s += radCoef[tapBase + t] * elemHist[(u32(jj) % D) * E + e]; }
246 }
247 return s;
248}
250@compute @workgroup_size(${GATHER_ELEMS})
251fn gatherPass(@builtin(local_invocation_id) lid: vec3<u32>, @builtin(workgroup_id) wid: vec3<u32>) {
252 // one thread per element; each workgroup writes its partial sums, which the host adds up
253 let t = lid.x; let e = wid.x * ${GATHER_ELEMS}u + t;
254 let step = bitcast<u32>(syncG[0]); // the node pass of step n has run: the finished sample belongs to s = n - 1
255 let j = i32((step - 2u) / M) - i32(W / 2u);
256 var fs = 0.0; var bs = 0.0;
257 if (j >= 0 && e < E) {
258 let hrow = OFF_NHIST + (u32(j) % D) * N;
259 let mean = (radR[hrow + elemUG[3u * e]] + radR[hrow + elemUG[3u * e + 1u]] + radR[hrow + elemUG[3u * e + 2u]]) / 3.0;
260 elemHist[(u32(j) % D) * E + e] = mean;
261 let rc = RC * e;
262 fs = radCoef[rc] * delayed(e, j, i32(radCoef[rc + 2u]), rc + 4u);
263 bs = radCoef[rc + 1u] * delayed(e, j, i32(radCoef[rc + 3u]), rc + 4u + T);
264 }
265 redF[t] = fs; redB[t] = bs;
266 workgroupBarrier();
267 for (var s = ${GATHER_ELEMS / 2}u; s > 0u; s = s >> 1u) {
268 if (t < s) { redF[t] += redF[t + s]; redB[t] += redB[t + s]; }
269 workgroupBarrier();
270 }
271 if (t == 0u && j >= 0) {
272 let idx = j - chunk.base;
273 if (idx >= 0 && u32(idx) < chunk.size) { pressure[u32(idx) * GE + wid.x] = redF[0]; pressure[(chunk.size + u32(idx)) * GE + wid.x] = redB[0]; }
274 }
275}
277// ------------------------------------------------------------------ capture pass (displacement frames)
278@group(3) @binding(0) var<storage, read> stateC: array<f32>;
279@group(3) @binding(1) var<storage, read_write> movie: array<f32>;
280@group(3) @binding(2) var<uniform> chunkC: Chunk;
281@group(3) @binding(3) var<storage, read> syncC: array<f32>;
283@compute @workgroup_size(256)
284fn capturePass(@builtin(global_invocation_id) gid: vec3<u32>) {
285 let nd = gid.x;
286 if (nd >= N) { return; }
287 let step = bitcast<u32>(syncC[0]);
288 let frame = (step / M) / chunkC.every - chunkC.frameBase;
289 movie[frame * N + nd] = stateC[5u * nd + 2u];
290}
291`;
292}