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