src/gpu/shaders.ts at main - proof-of-concept/cymsimweb

proof-of-concept / cymsimweb

328 lines · 16.5 KBBlame History Raw

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 /** Inverse of the M-Gram matrix of the six raw rigid modes (rawRigidModes in shell.ts), row-major. */

14 rigidGramInv: ArrayLike<number>;

15}

17// Per-element and per-node data are stored field-major (field f of element e at f * E + e), so that

18// neighbouring threads read neighbouring words.

19export const ELEM_STRIDE = 58; // fields: Gx[3] Gy[3] A z0x z0y C11 C12 C33 cd Ke[45] (upper triangle, row by row)

20export const NODE_STRIDE = 16; // fields: Minv[5] M[5] X Y z0 normal[3] (X, Y from the rigid centre); the decimation taps follow the nodes

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

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

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

24/** sync[2] when the partials are raw-mode sums; states exported before step 3 hold orthonormal-mode sums there and 0 here. */

25export const RAW_BASIS = 1;

26export const RADCOEF_STRIDE = 36; // coefF coefB i0F i0B tapsF[16] tapsB[16]

27export const WORKGROUP_NODES = 256;

28export const GATHER_ELEMS = 256; // elements per gather workgroup

29export const STATE_FLOATS = 13; // per node: q[5] v[5] vPrev[3]

30/** Chunk uniform: base size frameBase every accBase accRows pad pad (eight 32-bit words). */

31export const CHUNK_WORDS = 8;

33function f(x: number): string {

34 const s = x.toPrecision(9);

35 return /[.e]/.test(s) ? s : s + ".0";

36}

38/** -K_e U added to the bending forces, from the 45 stored entries of the symmetric K_e, written out

39 * in full: a loop over the packed triangle indexes private arrays dynamically, which was slower. */

40function bendingProduct(): string {

41 const lines: string[] = [], terms: string[][] = Array.from({ length: 9 }, () => []);

42 let f = 13;

43 for (let i = 0; i < 9; i++) for (let j = i; j < 9; j++) {

44 lines.push(` let k${i}${j} = elemF[${f++}u * E + e];`);

45 terms[i].push(`k${i}${j} * U[${j}]`);

46 if (j !== i) terms[j].push(`k${i}${j} * U[${i}]`);

47 }

48 for (let i = 0; i < 9; i++) lines.push(` out[${5 * Math.floor(i / 3) + 2 + (i % 3)}u] -= ${terms[i].join(" + ")};`);

49 return lines.join("\n");

50}

52export function shaderSource(c: ShaderConstants): string {

53 const HF = HIT_HEADER + 5 * c.P, HU = HITU_HEADER + 3 * c.P;

54 return /* wgsl */ `

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

56const 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;

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

58const HF: u32 = ${HF}u; const HU: u32 = ${HU}u; const RC: u32 = ${RADCOEF_STRIDE}u;

59const OFF_V: u32 = 5u * N; const OFF_VWP: u32 = 10u * N;

60const OFF_NHIST: u32 = W * N; const OFF_AHIST: u32 = (W + D) * N; const OFF_ACCOUT: u32 = (W + D + M) * N;

61const GINV = array<f32, 36>(${Array.from(c.rigidGramInv, f).join(", ")});

62const GE: u32 = ${Math.ceil(c.E / GATHER_ELEMS)}u;

63const OFF_TAPS: u32 = ${NODE_STRIDE}u * N; const OFF_PART: u32 = ${SYNC_HEADER}u;

65struct Chunk { base: i32, size: u32, frameBase: u32, every: u32, accBase: i32, accRows: u32, pad0: u32, pad1: u32 }

67// ------------------------------------------------------------------ element pass

68@group(0) @binding(0) var<storage, read> elemF: array<f32>;

69@group(0) @binding(1) var<storage, read> elemU: array<u32>;

70@group(0) @binding(2) var<storage, read> stateR: array<f32>;

71@group(0) @binding(3) var<storage, read_write> elemForce: array<f32>;

72@group(0) @binding(4) var<storage, read_write> hitsF: array<f32>;

73@group(0) @binding(5) var<storage, read> hitsUR: array<u32>;

74@group(0) @binding(6) var<storage, read_write> syncE: array<f32>;

76@compute @workgroup_size(64)

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

78 let id = gid.x;

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

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

81 if (id < E) {

82 let e = id;

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

84 var qe: array<f32, 15>;

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

86 let gx = vec3<f32>(elemF[e], elemF[1u * E + e], elemF[2u * E + e]);

87 let gy = vec3<f32>(elemF[3u * E + e], elemF[4u * E + e], elemF[5u * E + e]);

88 let A = elemF[6u * E + e]; let z0x = elemF[7u * E + e]; let z0y = elemF[8u * E + e];

89 let C11 = elemF[9u * E + e]; let C12 = elemF[10u * E + e]; let C33 = elemF[11u * E + e]; let cd = elemF[12u * E + e];

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

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

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

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

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

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

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

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

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

99 var out: array<f32, 15>;

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

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

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

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

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

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

106 }

107 var U: array<f32, 9>;

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

109${bendingProduct()}

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

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

112 // contact forces of hit k at this step

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

114 var wp: array<f32, ${c.P}>;

115 var wmax = -1.0e30;

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

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

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

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

120 }

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

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

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

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

125 var F = 0.0;

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

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

128 var fp = 0.0;

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

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

131 }

132 hitsF[hf + 7u] = F;

133 }

134}

135

136// ------------------------------------------------------------------ node pass

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

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

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

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

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

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

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

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

145@group(1) @binding(8) var<uniform> chunkN: Chunk;

146

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

148var<workgroup> csum: array<f32, 12>;

149var<workgroup> coef: array<f32, 12>;

150

151fn modW(x: i32) -> u32 { return u32(((x % i32(W)) + i32(W)) % i32(W)); }

152

153@compute @workgroup_size(${WORKGROUP_NODES})

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

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

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

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

158 // twelve threads sum the previous step's partials of all workgroups, p = R^T M v and R^T M q with R

159 // the raw rigid modes, and form the coefficients G^-1 p of the M-orthogonal projection; the others

160 // read the coefficients from shared memory

161 if (t < 12u) {

162 var s = 0.0;

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

164 csum[t] = s;

165 }

166 workgroupBarrier();

167 if (t < 12u) {

168 let r = t % 6u; let o = t - r;

169 var s = 0.0;

170 for (var m = 0u; m < 6u; m++) { s += GINV[6u * r + m] * csum[o + m]; }

171 coef[t] = s;

172 }

173 workgroupBarrier();

174 var c: array<f32, 12>;

175 for (var k = 0u; k < 12u; k++) { c[k] = coef[k]; }

176 var part: array<f32, 12>;

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

178 if (nd < N) {

179 var fo: array<f32, 5>;

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

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

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

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

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

185 }

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

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

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

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

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

191 }

192 }

193 }

194 }

195 var acc = 0.0;

196 // the rigid content to remove, sum_k c_k R_k at this node (the modes of rawRigidModes)

197 let X = nodeF[10u * N + nd]; let Y = nodeF[11u * N + nd]; let Z = nodeF[12u * N + nd];

198 let cv = array<f32, 5>(c[0] - c[2] * Y + c[5] * Z, c[1] + c[2] * X - c[4] * Z, c[3] + c[4] * Y - c[5] * X, c[4], c[5]);

199 let cq = array<f32, 5>(c[6] - c[8] * Y + c[11] * Z, c[7] + c[8] * X - c[10] * Z, c[9] + c[10] * Y - c[11] * X, c[10], c[11]);

200 var mv: array<f32, 5>; var mq: array<f32, 5>;

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

202 let idx = 5u * nd + d;

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

204 vd -= cv[d]; qd -= cq[d];

205 let v1 = (vd + DT * fo[d] * nodeF[d * N + nd]) * DAMP0;

206 // normal acceleration of the previous step from consecutive projected velocities

207 if (d < 3u) { acc += nodeF[(13u + d) * N + nd] * (vd - state[OFF_VWP + 3u * nd + d]) / DT; state[OFF_VWP + 3u * nd + d] = vd; }

208 let q1 = qd + DT * v1;

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

210 let mm = nodeF[(5u + d) * N + nd];

211 mv[d] = mm * v1; mq[d] = mm * q1;

212 }

213 // this node's share of R^T M v' and R^T M q'

214 part[0] = mv[0]; part[1] = mv[1]; part[2] = -Y * mv[0] + X * mv[1]; part[3] = mv[2]; part[4] = -Z * mv[1] + Y * mv[2] + mv[3]; part[5] = Z * mv[0] - X * mv[2] + mv[4];

215 part[6] = mq[0]; part[7] = mq[1]; part[8] = -Y * mq[0] + X * mq[1]; part[9] = mq[2]; part[10] = -Z * mq[1] + Y * mq[2] + mq[3]; part[11] = Z * mq[0] - X * mq[2] + mq[4];

216 // decimation of the acceleration of sub-step s = step - 1 into the W accumulator slots of this node;

217 // slot w belongs to audio sample j0 + w - W/2 + 1, kept at ring position (sample mod W)

218 if (step >= 1u) {

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

220 if (kk == 0u) {

221 let jE = i32(j0) - i32(W / 2u);

222 let slotE = modW(jE);

223 let val = rad[slotE * N + nd];

224 if (jE >= 0) {

225 rad[OFF_NHIST + (u32(jE) % D) * N + nd] = val;

226 let row = jE - chunkN.accBase;

227 if (chunkN.accRows > 0u && row >= 0 && u32(row) < chunkN.accRows) { rad[OFF_ACCOUT + u32(row) * N + nd] = val; }