Blame src/gpu/shaders.ts at 83544826aa7a1fd2b69c2a355f975d0a9cc5077f

8354482 cymmodes: linear modes of a cymsim cymbal driven in the full nonlinear model, with the energy of every mode over timeJeremy Magland 19 days ago1// WGSL kernels of the time loop, as in cymsimweb. Entry points: the element pass (forces, sigma1

2// term, contact forces of the hits, decimation accumulators of the radiated acceleration), the

3// node pass (gather, contact, distributed drive, lagged rigid-mode projection, leapfrog update,

4// partial sums, exciter update), the gather pass (retarded-time sum at the listener once per

5// audio sample), the capture pass (displacement movie) and, new here, the projection pass, which

6// contracts the transverse displacement and velocity with a mass-weighted mode matrix and writes

7// one modal snapshot. Sizes and constants are baked into the source; K = 0 (no hits) is allowed.

9export interface ShaderConstants {

10 N: number; E: number; K: number; P: number; M: number; W: number; D: number; G: number;

11 dt: number; sigma0: number; nl: number;

12 Nm: number; // modes in the projection (1 when unused)

13 driveOmega: number; driveRampSteps: number; driveStopStep: number; // stopStep 0 disables

14 traceMode: number; // mode traced at the audio rate; -1 disables

15}

17export const ELEM_STRIDE = 96; // Gx[3] Gy[3] A z0x z0y C11 C12 C33 cd Ke[81] pad[2]

18export const NODE_STRIDE = 40; // Minv[5] M[5] Rm[6][5]

19export const HIT_HEADER = 12; // kPoint alpha mass gap velocity zs vs F Fmax pad[3]

20export const HITU_HEADER = 4; // releaseStep ncont pad[2]

21export const SYNC_HEADER = 4; // stepA stepB pad[2], then partials[2][G][12]

22export const WORKGROUP_NODES = 256;

23export const CHUNK_UNIFORM_BYTES = 32;

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 const PA = Math.max(1, c.P); // array<f32, 0> is invalid WGSL; the hit branch never runs when K = 0

33 return /* wgsl */ `

34const N: u32 = ${c.N}u; const E: u32 = ${c.E}u; const K: u32 = ${c.K}u; const P: u32 = ${c.P}u;

35const M: u32 = ${c.M}u; const W: u32 = ${c.W}u; const D: u32 = ${c.D}u; const G: u32 = ${c.G}u;

36const DT: f32 = ${f(c.dt)}; const DAMP0: f32 = ${f(1 / (1 + 2 * c.sigma0 * c.dt))}; const NL: f32 = ${f(c.nl)};

37const EF: u32 = ${ELEM_STRIDE}u; const NF: u32 = ${NODE_STRIDE}u; const HF: u32 = ${HF}u; const HU: u32 = ${HU}u;

38const NM: u32 = ${Math.max(1, c.Nm)}u; const TRACE_MODE: u32 = ${Math.max(0, c.traceMode)}u;

39const DRV_STOP: u32 = ${Math.max(0, Math.floor(c.driveStopStep))}u;

40const DRV_OMDT: f32 = ${f(c.driveOmega * c.dt)}; const DRV_RAMP: f32 = ${f(Math.max(1, c.driveRampSteps))};

41const OFF_V: u32 = 5u * N; const OFF_ACC: u32 = 10u * N; const OFF_VWP: u32 = 11u * N;

42const OFF_HIST: u32 = E * W; const OFF_TAPS: u32 = EF * E; const OFF_PART: u32 = ${SYNC_HEADER}u;

44struct Chunk { base: i32, size: u32, frameBase: u32, every: u32, projBase: u32, projEvery: u32, phase0: f32, step0: 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>;

54@group(0) @binding(7) var<storage, read_write> rad: array<f32>;

56@compute @workgroup_size(64)

57fn elementPass(@builtin(global_invocation_id) gid: vec3<u32>) {

58 let id = gid.x;

59 let step = bitcast<u32>(syncE[0]);

60 if (id == 0u) { syncE[1] = syncE[0]; }

61 if (id < E) {

62 let e = id;

63 let b = e * EF;

64 let n0 = elemU[3u * e]; let n1 = elemU[3u * e + 1u]; let n2 = elemU[3u * e + 2u];

65 var qe: array<f32, 15>;

66 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]; }

67 let gx = vec3<f32>(elemF[b], elemF[b + 1u], elemF[b + 2u]);

68 let gy = vec3<f32>(elemF[b + 3u], elemF[b + 4u], elemF[b + 5u]);

69 let A = elemF[b + 6u]; let z0x = elemF[b + 7u]; let z0y = elemF[b + 8u];

70 let C11 = elemF[b + 9u]; let C12 = elemF[b + 10u]; let C33 = elemF[b + 11u]; let cd = elemF[b + 12u];

71 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]);

72 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);

73 let ax = NL * wx + z0x; let ay = NL * wy + z0y;

74 let exx = ux + NL * 0.5 * wx * wx + z0x * wx;

75 let eyy = vy + NL * 0.5 * wy * wy + z0y * wy;

76 let gxy = uy + vx + NL * wx * wy + z0x * wy + z0y * wx;

77 let Nxx = C11 * exx + C12 * eyy; let Nyy = C12 * exx + C11 * eyy; let Nxy = C33 * gxy;

78 let vw = vec3<f32>(stateR[OFF_V + 5u * n0 + 2u], stateR[OFF_V + 5u * n1 + 2u], stateR[OFF_V + 5u * n2 + 2u]);

79 let vwx = dot(gx, vw); let vwy = dot(gy, vw);

80 var out: array<f32, 15>;

81 for (var i = 0u; i < 3u; i++) {

82 let g_x = gx[i]; let g_y = gy[i];

83 out[5u * i] = -A * (Nxx * g_x + Nxy * g_y);

84 out[5u * i + 1u] = -A * (Nyy * g_y + Nxy * g_x);

85 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);

86 out[5u * i + 3u] = 0.0; out[5u * i + 4u] = 0.0;

87 }

88 var U: array<f32, 9>;

89 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]; }

90 for (var i = 0u; i < 9u; i++) {

91 var s = 0.0;

92 for (var j = 0u; j < 9u; j++) { s += elemF[b + 13u + 9u * i + j] * U[j]; }

93 let slot = 5u * (i / 3u) + 2u + (i % 3u);

94 out[slot] -= s;

95 }

96 for (var i = 0u; i < 15u; i++) { elemForce[15u * e + i] = out[i]; }

97 // decimation accumulators of the radiated acceleration (of the previous step)

98 if (step >= 1u) {

99 let s = step - 1u; let j0 = s / M; let kk = s - j0 * M;

100 if (kk == 0u && j0 > 0u) {

101 let j = i32(j0) - i32(W / 2u);

102 if (j >= 0) { rad[OFF_HIST + e * D + (u32(j) % D)] = rad[e * W]; }

103 for (var w = 0u; w < W - 1u; w++) { rad[e * W + w] = rad[e * W + w + 1u]; }

104 rad[e * W + W - 1u] = 0.0;

105 }

106 let a = (stateR[OFF_ACC + n0] + stateR[OFF_ACC + n1] + stateR[OFF_ACC + n2]) / 3.0;

107 for (var w = 0u; w < W; w++) { rad[e * W + w] += a * elemF[OFF_TAPS + w * M + kk]; }

108 }

109 } else if (id < E + K) {

110 // contact forces of hit k at this step

111 let k = id - E; let hf = k * HF; let hu = k * HU;

112 var wp: array<f32, ${PA}>;

113 var wmax = -1.0e30;

114 for (var p = 0u; p < P; p++) {

115 var w = hitsF[hf + ${HIT_HEADER}u + 3u * P + p];

116 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]; }

117 wp[p] = w; wmax = max(wmax, w);

118 }

119 var zs = hitsF[hf + 5u];

120 if (step == hitsUR[hu]) { zs = wmax + hitsF[hf + 3u]; hitsF[hf + 5u] = zs; hitsF[hf + 6u] = -hitsF[hf + 4u]; }

121 let isActive = step >= hitsUR[hu];

122 let kPoint = hitsF[hf]; let alpha = hitsF[hf + 1u];

123 var F = 0.0;

124 for (var p = 0u; p < P; p++) {

125 let delta = max(wp[p] - zs, 0.0);

126 var fp = 0.0;

127 if (isActive && delta > 0.0) { fp = kPoint * pow(delta, alpha); }

128 hitsF[hf + ${HIT_HEADER}u + 4u * P + p] = fp; F += fp;

129 }

130 hitsF[hf + 7u] = F;

131 }

132}

133

134// ------------------------------------------------------------------ node pass

135@group(1) @binding(0) var<storage, read_write> state: array<f32>;

136@group(1) @binding(1) var<storage, read> elemForceR: array<f32>;

137@group(1) @binding(2) var<storage, read> nodeF: array<f32>;

138@group(1) @binding(3) var<storage, read> adj: array<u32>;

139@group(1) @binding(4) var<storage, read_write> hitsFN: array<f32>;

140@group(1) @binding(5) var<storage, read_write> hitsU: array<u32>;

141@group(1) @binding(6) var<storage, read_write> sync: array<f32>;

142@group(1) @binding(7) var<storage, read> driveF: array<f32>;

143@group(1) @binding(8) var<uniform> chunkD: Chunk;

144

145var<workgroup> red: array<f32, ${12 * WORKGROUP_NODES}>;

146

147@compute @workgroup_size(${WORKGROUP_NODES})

148fn nodePass(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>, @builtin(workgroup_id) wid: vec3<u32>) {

149 let nd = gid.x; let t = lid.x;

150 let step = bitcast<u32>(sync[1]);

151 let parity = step & 1u; let prevp = (step + 1u) & 1u;

152 var c: array<f32, 12>;

153 for (var k = 0u; k < 12u; k++) {

154 var s = 0.0;

155 for (var g = 0u; g < G; g++) { s += sync[OFF_PART + (prevp * G + g) * 12u + k]; }

156 c[k] = s;

157 }

158 var part: array<f32, 12>;

159 for (var k = 0u; k < 12u; k++) { part[k] = 0.0; }

160 if (nd < N) {

161 var fo: array<f32, 5>;

162 for (var d = 0u; d < 5u; d++) { fo[d] = 0.0; }

163 let a0 = adj[nd]; let a1 = adj[nd + 1u];

164 for (var a = a0; a < a1; a++) {

165 let ent = adj[N + 1u + a]; let e = ent / 3u; let i = ent - 3u * e;

166 for (var d = 0u; d < 5u; d++) { fo[d] += elemForceR[15u * e + 5u * i + d]; }

167 }

168 for (var k = 0u; k < K; k++) {

169 for (var p = 0u; p < P; p++) {

170 for (var i = 0u; i < 3u; i++) {

171 if (hitsU[k * HU + ${HITU_HEADER}u + 3u * p + i] == nd) {

172 fo[2] -= hitsFN[k * HF + ${HIT_HEADER}u + 4u * P + p] * hitsFN[k * HF + ${HIT_HEADER}u + 3u * p + i];

173 }

174 }

175 }

176 }

177 // distributed drive on w: fNode * ramp * sin(omega t), phase carried across chunks in phase0

178 if (DRV_STOP > 0u && step < DRV_STOP) {

179 let arg = chunkD.phase0 + f32(step - chunkD.step0) * DRV_OMDT;

180 let env = min(f32(step) / DRV_RAMP, 1.0);

181 fo[2] += driveF[nd] * env * sin(arg);

182 }

183 let nb = nd * NF;

184 for (var d = 0u; d < 5u; d++) {

185 let idx = 5u * nd + d;

186 var vd = state[OFF_V + idx]; var qd = state[idx];

187 for (var k = 0u; k < 6u; k++) { let r = nodeF[nb + 10u + 5u * k + d]; vd -= c[k] * r; qd -= c[6u + k] * r; }

188 let v1 = (vd + DT * fo[d] * nodeF[nb + d]) * DAMP0;

189 if (d == 2u) { state[OFF_ACC + nd] = (vd - state[OFF_VWP + nd]) / DT; state[OFF_VWP + nd] = vd; }

190 let q1 = qd + DT * v1;

191 state[idx] = q1; state[OFF_V + idx] = v1;

192 let mm = nodeF[nb + 5u + d];

193 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; }

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}

215

216// ------------------------------------------------------------------ gather pass (listener)

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

223var<workgroup> redF: array<f32, 256>;

224var<workgroup> redB: array<f32, 256>;

225

226@compute @workgroup_size(256)

227fn gatherPass(@builtin(local_invocation_id) lid: vec3<u32>) {

228 let t = lid.x;

229 let step = bitcast<u32>(syncG[0]);

230 let j = i32((step - 1u) / M) - i32(W / 2u);

231 var fs = 0.0; var bs = 0.0;

232 if (j >= 0) {

233 for (var e = t; e < E; e += 256u) {

234 let rc = 6u * e;

235 let jf = j - i32(radCoef[rc + 2u]); let jb = j - i32(radCoef[rc + 4u]);

236 if (jf >= 1) {

237 let fr = radCoef[rc + 3u];

238 fs += radCoef[rc] * (radR[OFF_HIST + e * D + (u32(jf) % D)] * (1.0 - fr) + radR[OFF_HIST + e * D + (u32(jf - 1) % D)] * fr);

239 }

240 if (jb >= 1) {

241 let fr = radCoef[rc + 5u];

242 bs += radCoef[rc + 1u] * (radR[OFF_HIST + e * D + (u32(jb) % D)] * (1.0 - fr) + radR[OFF_HIST + e * D + (u32(jb - 1) % D)] * fr);

243 }

244 }

245 }

246 redF[t] = fs; redB[t] = bs;

247 workgroupBarrier();

248 for (var s = 128u; s > 0u; s = s >> 1u) {

249 if (t < s) { redF[t] += redF[t + s]; redB[t] += redB[t + s]; }

250 workgroupBarrier();

251 }

252 if (t == 0u && j >= 0) {

253 let idx = j - chunk.base;

254 if (idx >= 0 && u32(idx) < chunk.size) { pressure[u32(idx)] = redF[0]; pressure[chunk.size + u32(idx)] = redB[0]; }

255 }

256}

257

258// ------------------------------------------------------------------ capture pass (displacement movie)

259// and projection pass (modal snapshots); they share group 3, with disjoint extra bindings, and

260// each pipeline's layout covers only the bindings its entry point uses.

261@group(3) @binding(0) var<storage, read> stateC: array<f32>;

262@group(3) @binding(1) var<storage, read_write> movie: array<f32>;

263@group(3) @binding(2) var<uniform> chunkC: Chunk;

264@group(3) @binding(3) var<storage, read> syncC: array<f32>;

265@group(3) @binding(4) var<storage, read> phiM: array<f32>;

266@group(3) @binding(5) var<storage, read_write> projOut: array<f32>;

267@group(3) @binding(6) var<storage, read_write> traceOut: array<f32>;

268

269@compute @workgroup_size(256)

270fn capturePass(@builtin(global_invocation_id) gid: vec3<u32>) {

271 let nd = gid.x;

272 if (nd >= N) { return; }

273 let step = bitcast<u32>(syncC[0]);

274 let frame = (step / M) / chunkC.every - chunkC.frameBase;

275 movie[frame * N + nd] = stateC[5u * nd + 2u];

276}

277

278var<workgroup> redQ: array<f32, 256>;

279var<workgroup> redV: array<f32, 256>;

280

281@compute @workgroup_size(256)

282fn projPass(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {

283 let m = wid.x; let t = lid.x;

284 var sq = 0.0; var sv = 0.0;

285 for (var nd = t; nd < N; nd += 256u) {

286 let p = phiM[m * N + nd];

287 sq += p * stateC[5u * nd + 2u];

288 sv += p * stateC[OFF_V + 5u * nd + 2u];

289 }

290 redQ[t] = sq; redV[t] = sv;

291 workgroupBarrier();

292 for (var s = 128u; s > 0u; s = s >> 1u) {

293 if (t < s) { redQ[t] += redQ[t + s]; redV[t] += redV[t + s]; }

294 workgroupBarrier();

295 }

296 if (t == 0u) {

297 let step = bitcast<u32>(syncC[0]);

298 let snap = (step / M) / chunkC.projEvery - chunkC.projBase;

299 projOut[snap * 2u * NM + m] = redQ[0];

300 projOut[snap * 2u * NM + NM + m] = redV[0];

301 }

302}

303

304// ------------------------------------------------------------------ trace pass (one mode's

305// coordinate at the audio rate, one workgroup per audio sample)

306var<workgroup> redT: array<f32, 256>;

307

308@compute @workgroup_size(256)

309fn tracePass(@builtin(local_invocation_id) lid: vec3<u32>) {

310 let t = lid.x;

311 var sq = 0.0;

312 for (var nd = t; nd < N; nd += 256u) {

313 sq += phiM[TRACE_MODE * N + nd] * stateC[5u * nd + 2u];

314 }

315 redT[t] = sq;

316 workgroupBarrier();

317 for (var s = 128u; s > 0u; s = s >> 1u) {

318 if (t < s) { redT[t] += redT[t + s]; }

319 workgroupBarrier();

320 }

321 if (t == 0u) {

322 let step = bitcast<u32>(syncC[0]);

323 traceOut[step / M - chunkC.step0 / M] = redT[0];

324 }

325}

326`;

327}