1/**
2 * WebGPU port of plate_fd_loop.m: the explicit time stepping of the plate,
3 * run as compute shaders. Everything around the loop stays in numbl: the host
4 * runs the unmodified plate_fd_setup.m, hands its struct in here as FdLoopSetup,
5 * and passes the record this returns to the unmodified plate_fd_output.m, which
6 * computes the stresses, the audio and the spectra in double precision.
7 *
8 * One time step is five compute dispatches, the dispatch boundaries being the
9 * global synchronization points the stencils need:
10 * strain e = G phi: every difference stencil of the energy method (second
11 * differences at the nodes, twists at the cells, first differences
12 * along the edges) as an exact integer stencil applied to phi
13 * force at every node, G' applied to the weighted strains: the bending
14 * force, the prestress and spring force, and the Laplacian, each
15 * from its own weights. This is the assembled operator applied in
16 * its factored form. In single precision the assembled 13-point
17 * stencil loses the smooth modes to rounding (the biharmonic of a
18 * smooth field is a small residual of large terms, and the rounding
19 * of the field alone is percent of it at low wavenumbers); the
20 * factored form keeps them to about 1e-5, because a second
21 * difference of nearby values is nearly exact in f32 and the
22 * scales live in the weights, where rounding is harmless
23 * dots one workgroup: the nodal forces f (strike, contacts, dashpots) from
24 * the total field w = R c + phi and its velocity, the reductions
25 * R' f, (AR)' phi and the support sums, the rigid-body update
26 * c -> cn (thread 0), and the series of this step (probe, rigid
27 * coordinates, support forces) averaged into their samples
28 * update phin = c_new (2 phi - c_old phim + dt^2 (f - B - AR c)/M - dt^2 R cacc
29 * + 2 sigma1 dt (Lphi - Lphi_prev)), zeroed at nodes that are not
30 * unknowns; frames recorded
31 * reduce the acceleration of the total field summed per listener delay
32 * bin into the audio samples; max |phin| plus the rigid extent for
33 * the blow-up test
34 *
35 * Steps are recorded in chunks of a few thousand into one command buffer each,
36 * with the step index a chunk base (ctl[0], written by the host once per chunk)
37 * plus a within-chunk offset from a dynamic-offset uniform; every recorded
38 * dispatch is identical. Two chunks are in flight at a time and the blow-up
39 * flag is read back after each, so a diverging run stops a chunk late rather
40 * than at the end.
41 *
42 * Precision. The GPU works in single precision. The flexural field is carried
43 * as (phi, d = phi - phim) and the rigid coordinates as (c, dc): the update is
44 * algebraically the .m's,
45 * dn = (d + rest) - cs (2 d + rest), phin = phi + dn, cs = sigma0 dt / (1 + sigma0 dt),
46 * with rest = dt^2 acc (the sigma1 loss being part of the nodal force), so that no factor of
47 * the form 1 - epsilon is stored in f32 and the rounding of the field does not
48 * feed back into the velocity. The split into rigid coordinates and flexural
49 * remainder is the .m's own (plate_rigid_basis); bending is never applied to
50 * the rigid part, whose forces come from the precomputed AR and Krr.
51 */
53export interface Triplets {
54 /** 1-based row and column indices, as MATLAB's find returns them */
55 i: Float64Array;
56 j: Float64Array;
57 v: Float64Array;
58}
60/** The fields of plate_fd_setup's struct that the loop reads; indices are 1-based as in the .m. */
61export interface FdLoopSetup {
62 nx: number;
63 ny: number;
64 nsteps: number;
65 dx: number;
66 dy: number;
67 dt: number;
68 fs: number;
69 /** Number of output samples */
70 na: number;
71 keep: Float64Array;
72 Mvec: Float64Array;
73 /**
74 * The operators in factored form (plate_fd_setup): Gs stacks the difference
75 * stencils (ng rows over the N nodes), Gscale the scale of each row, and the
76 * symmetric weights give KB = Gs' WB Gs, KG - diag(ks) = Gs' WG Gs and
77 * Lm = diag(inva) Gs' WL Gs on the unknowns.
78 */
79 Gs: Triplets;
80 Gscale: Float64Array;
81 ng: number;
82 WB: Triplets;
83 WG: Triplets;
84 WL: Triplets;
85 inva: Float64Array;
86 /** Linear springs, the diagonal part of KG */
87 ks: Float64Array;
88 nr: number;
89 /** N x nr column-major */
90 R: Float64Array;
91 AR: Float64Array;
92 /** nr x nr column-major */
93 Krr: Float64Array;
94 Ginv: Float64Array;
95 kb: Float64Array;
96 cb: Float64Array;
97 gb: Float64Array;
98 kt: Float64Array;
99 ct: Float64Array;
100 gt: Float64Array;
101 cs: Float64Array;
102 /** Static force at rest of the springs, the bottom contacts and the top contacts */
103 fstatS: Float64Array;
104 fstatB: Float64Array;
105 fstatT: Float64Array;
106 /** Supports as a CSR list of 1-based node indices, each reporting one kind of force (0 spring, 1 bottom, 2 top) */
107 supPtr: Float64Array;
108 supNodes: Float64Array;
109 supKind: Float64Array;
110 nsup: number;
111 /** Force per node per unit F */
112 sa: Float64Array;
113 /** Force at the nsteps + 1 time levels */
114 F: Float64Array;
115 ip: number;
116 jp: number;
117 frameSteps: Float64Array;
118 /** Listener entries: node (1-based), delay bin (1-based), weight */
119 lisNode: Float64Array;
120 lisBin: Float64Array;
121 lisWgt: Float64Array;
122 nbins: number;
123 c0: Float64Array;
124 phi0: Float64Array;
125 sigma0: number;
126 /** The sigma1 loss as a nodal force per unit Lm(w - wm): 2 sigma1 M / dt, per node */
127 s1f: Float64Array;
128 /** Lm R, N x nr column-major: the Laplacian per node of the rigid basis */
129 LR: Float64Array;
130 dt2: number;
131 /** Weight of a step in its output sample (1 over the steps the sample contains), na entries */
132 wsamp: Float64Array;
133 blowup: number;
134}
136/** What plate_fd_loop.m returns, in single precision, before conversion for plate_fd_output.m. */
137export interface FdLoopRecord {
138 w_p: Float32Array;
139 wxx_p: Float32Array;
140 wyy_p: Float32Array;
141 wxy_p: Float32Array;
142 /** nr x na, row k at C[k * na ..] */
143 C: Float32Array;
144 /** nsup x na */
145 Fsup: Float32Array;
146 /** ny x nx x nf */
147 Wf: Float32Array;
148 /** nr x nf, row k at Cf[k * nf ..] */
149 Cf: Float32Array;
150 /** nbins x na */
151 Pbins: Float32Array;
152 divergedStep: number;
153 gpuMs: number;
154}
156export interface GpuHandle {
157 device: GPUDevice;
158 adapter: string;
159}
161export function describeAdapter(adapter: GPUAdapter): string {
162 const info = (adapter as { info?: Partial<GPUAdapterInfo> }).info;
163 if (!info) return "unknown adapter";
164 const name = [info.vendor, info.architecture].filter(Boolean).join(" ");
165 return name || info.description || "unknown adapter";
166}
168/** The device request the browser worker and the Node test both make. */
169export async function requestGpuDevice(): Promise<GpuHandle> {
170 const gpu = (globalThis as { navigator?: { gpu?: GPU } }).navigator?.gpu;
171 if (!gpu) throw new Error("WebGPU is not available");
172 const adapter = await gpu.requestAdapter();
173 if (!adapter) throw new Error("no WebGPU adapter");
174 const device = await adapter.requestDevice({
175 requiredLimits: {
176 maxStorageBufferBindingSize: adapter.limits.maxStorageBufferBindingSize,
177 maxBufferSize: adapter.limits.maxBufferSize,
178 },
179 });
180 return { device, adapter: describeAdapter(adapter) };
181}
183const WG = 64;
184const RWG = 256;
185const CHUNK = 2048;
186const NONE = 0xffffffff;
188const fmt = (x: number): string => {
189 if (!Number.isFinite(x)) throw new Error("non-finite kernel constant");
190 let s = x.toPrecision(9);
191 if (!/[.e]/.test(s)) s += ".0";
192 return s + "f";
193};
195interface Csr {
196 ptr: Uint32Array;
197 col: Uint32Array;
198 val: Float32Array;
199}
201function toCsr(t: Triplets, n: number): Csr {
202 const count = new Uint32Array(n);
203 for (let p = 0; p < t.i.length; p++) count[t.i[p] - 1]++;
204 const ptr = new Uint32Array(n + 1);
205 for (let r = 0; r < n; r++) ptr[r + 1] = ptr[r] + count[r];
206 const fill = ptr.slice(0, n);
207 const col = new Uint32Array(t.i.length);
208 const val = new Float32Array(t.i.length);
209 for (let p = 0; p < t.i.length; p++) {
210 const q = fill[t.i[p] - 1]++;
211 col[q] = t.j[p] - 1;
212 val[q] = t.v[p];
213 }
214 return { ptr, col, val };
215}
217/** Sparse product in double. */
218function spmv(t: Triplets, x: Float64Array, n: number): Float64Array {
219 const y = new Float64Array(n);
220 for (let p = 0; p < t.i.length; p++) y[t.i[p] - 1] += t.v[p] * x[t.j[p] - 1];
221 return y;
222}
224/** Transposed sparse product in double. */
225function spmvT(t: Triplets, x: Float64Array, n: number): Float64Array {
226 const y = new Float64Array(n);
227 for (let p = 0; p < t.i.length; p++) y[t.j[p] - 1] += t.v[p] * x[t.i[p] - 1];
228 return y;
229}
231/** The weight matrices as one CSR pattern with three value arrays. */
232function weightsCsr(
233 ng: number,
234 scale: Float64Array,
235 WB: Triplets,
236 WG: Triplets,
237 WL: Triplets,
238): { ptr: Uint32Array; col: Uint32Array; vb: Float32Array; vg: Float32Array; vl: Float32Array } {
239 const rows: Map<number, { b: number; g: number; l: number }>[] = Array.from(
240 { length: ng },
241 () => new Map(),
242 );
243 const add = (t: Triplets, key: "b" | "g" | "l") => {
244 for (let p = 0; p < t.i.length; p++) {
245 const r = t.i[p] - 1;
246 const c = t.j[p] - 1;
247 let e = rows[r].get(c);
248 if (!e) {
249 e = { b: 0, g: 0, l: 0 };
250 rows[r].set(c, e);
251 }
252 // the stencils are integers: their scales move into the weights
253 e[key] += t.v[p] * scale[r] * scale[c];
254 }
255 };
256 add(WB, "b");
257 add(WG, "g");
258 add(WL, "l");
259 const ptr = new Uint32Array(ng + 1);
260 for (let r = 0; r < ng; r++) ptr[r + 1] = ptr[r] + rows[r].size;
261 const nnz = ptr[ng];
262 const col = new Uint32Array(nnz);
263 const vb = new Float32Array(nnz);
264 const vg = new Float32Array(nnz);
265 const vl = new Float32Array(nnz);
266 let q = 0;
267 for (let r = 0; r < ng; r++)
268 for (const [c, e] of rows[r]) {
269 col[q] = c;
270 vb[q] = e.b;
271 vg[q] = e.g;
272 vl[q] = e.l;
273 q++;
274 }
275 return { ptr, col, vb, vg, vl };
276}
278/** Listener entries grouped by delay bin: entry -> node, entry -> weight, in bin order. */
279function binLists(
280 bins: Float64Array,
281 nodes: Float64Array,
282 wgt: Float64Array,
283 nbins: number,
284): { ptr: Uint32Array; node: Uint32Array; wgt: Float32Array } {
285 const n = bins.length;
286 const count = new Uint32Array(nbins);
287 for (let e = 0; e < n; e++) count[bins[e] - 1]++;
288 const ptr = new Uint32Array(nbins + 1);
289 for (let b = 0; b < nbins; b++) ptr[b + 1] = ptr[b] + count[b];
290 const fill = ptr.slice(0, nbins);
291 const node = new Uint32Array(n);
292 const w = new Float32Array(n);
293 for (let e = 0; e < n; e++) {
294 const q = fill[bins[e] - 1]++;
295 node[q] = nodes[e] - 1;
296 w[q] = wgt[e];
297 }
298 return { ptr, node, wgt: w };
299}
301function pack<T extends Uint32Array | Float32Array>(
302 parts: [string, ArrayLike<number>][],
303 make: (n: number) => T,
304): { data: T; off: Record<string, number> } {
305 const off: Record<string, number> = {};
306 let cursor = 0;
307 for (const [name, arr] of parts) {
308 off[name] = cursor;
309 cursor += arr.length;
310 }
311 const data = make(Math.max(1, cursor));
312 for (const [name, arr] of parts) data.set(arr, off[name]);
313 return { data, off };
314}
316function buildShaders(su: FdLoopSetup) {
317 const nx = su.nx | 0;
318 const ny = su.ny | 0;
319 const N = nx * ny;
320 const nsteps = su.nsteps | 0;
321 const nbins = su.nbins | 0;
322 const na = su.na | 0;
323 const nf = su.frameSteps.length;
324 const nr = su.nr | 0;
325 const nsup = su.nsup | 0;
326 if (nr !== 0 && nr !== 3) throw new Error("the loop expects 0 or 3 rigid-body coordinates");
327 if (su.sa.length !== N || su.keep.length !== N || su.Mvec.length !== N)
328 throw new Error("setup arrays do not match the grid");
329 if (su.F.length < nsteps) throw new Error("force history shorter than the step count");
330 if (su.wsamp.length !== na) throw new Error("sample weights do not match the sample count");
331 if (nr > 0 && (su.R.length !== N * nr || su.AR.length !== N * nr || su.LR.length !== N * nr))
332 throw new Error("rigid-body arrays do not match the grid");
333 if (su.s1f.length !== N) throw new Error("loss array does not match the grid");
334 if (nsup > 8) throw new Error("at most 8 supports");
336 // ---- per-step tables, decided in f64 exactly as the .m decides them
337 const frameOf = new Uint32Array(nsteps).fill(NONE);
338 for (let kf = 0; kf < nf; kf++) {
339 const step = su.frameSteps[kf];
340 if (step >= 1 && step <= nsteps) frameOf[step - 1] = kf;
341 }
342 const iaOf = new Uint32Array(nsteps);
343 for (let k = 0; k < nsteps; k++) {
344 const ia = Math.floor(k * su.dt * su.fs); // ia = floor((nstep - 1) dt fs) + 1, made 0-based
345 if (ia >= na) throw new Error("output sample index out of range");
346 iaOf[k] = ia;
347 }
348 const bl = binLists(su.lisBin, su.lisNode, su.lisWgt, nbins);
349 const ng = su.ng | 0;
350 if (su.Gscale.length !== ng) throw new Error("strain scales do not match the strain count");
351 // the stencils as exact integers (their scale is folded into the weights), and the transpose
352 const Gi: Triplets = {
353 i: su.Gs.i,
354 j: su.Gs.j,
355 v: Float64Array.from(su.Gs.v, (v, p) => Math.round(v / su.Gscale[su.Gs.i[p] - 1])),
356 };
357 for (let p = 0; p < Gi.v.length; p++)
358 if (Math.abs(Gi.v[p] - su.Gs.v[p] / su.Gscale[su.Gs.i[p] - 1]) > 1e-6)
359 throw new Error("a stencil coefficient is not an integer");
360 const G = toCsr(Gi, ng);
361 const GT = toCsr({ i: Gi.j, j: Gi.i, v: Gi.v }, N);
362 const W = weightsCsr(ng, su.Gscale, su.WB, su.WG, su.WL);
363 const minv = Float32Array.from(su.Mvec, (m) => 1 / m);
364 const f32 = (a: ArrayLike<number>) => Float32Array.from(a);
365 const rmax = new Float32Array(3);
366 for (let k = 0; k < nr; k++) {
367 let m = 0;
368 for (let i = 0; i < N; i++) m = Math.max(m, Math.abs(su.R[i + k * N]));
369 rmax[k] = m;
370 }
371 const supPtr = Uint32Array.from(su.supPtr, (p) => p - 1);
372 const supNodes = Uint32Array.from(su.supNodes, (p) => p - 1);
373 const supKind = Uint32Array.from(su.supKind);
375 const opsParts: [string, ArrayLike<number>][] = [
376 ["SA", f32(su.sa)],
377 ["F", f32(su.F.subarray(0, nsteps))],
378 ["MINV", minv],
379 ["KEEP", f32(su.keep)],
380 ["INVA", f32(su.inva)],
381 ["KSD", f32(su.ks)],
382 ["GV", G.val],
383 ["GTV", GT.val],
384 ["WBV", W.vb],
385 ["WGV", W.vg],
386 ["WLV", W.vl],
387 ["R", f32(nr ? su.R : [0])],
388 ["AR", f32(nr ? su.AR : [0])],
389 ["LR", f32(nr ? su.LR : [0])],
390 ["S1F", f32(su.s1f)],
391 ["KRR", f32(nr ? su.Krr : [0])],
392 ["GINV", f32(nr ? su.Ginv : [0])],
393 ["KB", f32(su.kb)],
394 ["CB", f32(su.cb)],
395 ["GB", f32(su.gb)],
396 ["KT", f32(su.kt)],
397 ["CT", f32(su.ct)],
398 ["GT", f32(su.gt)],
399 ["CS", f32(su.cs)],
400 ["KS", f32(su.ks)],
401 ["FSTATS", f32(su.fstatS)],
402 ["FSTATB", f32(su.fstatB)],
403 ["FSTATT", f32(su.fstatT)],
404 ["WGT", bl.wgt],
405 ["WSAMP", f32(su.wsamp)],
406 ];
407 const idxParts: [string, ArrayLike<number>][] = [
408 ["GPTR", G.ptr],
409 ["GCOL", G.col],
410 ["GTPTR", GT.ptr],
411 ["GTCOL", GT.col],
412 ["WPTR", W.ptr],
413 ["WCOL", W.col],
414 ["BINPTR", bl.ptr],
415 ["BINNODE", bl.node],
416 ["SUPPTR", supPtr],
417 ["SUPNODES", supNodes],
418 ["SUPKIND", supKind],
419 ["FRAME", frameOf],
420 ["IA", iaOf],
421 ];
422 const ops = pack(opsParts, (n) => new Float32Array(n));
423 const idx = pack(idxParts, (n) => new Uint32Array(n));
425 // ---- state buffer: two slots each of phi, d = phi - phim and Lphi; B; f; fsup; rigid scalars
426 const ST_PHI = 0;
427 const ST_D = 2 * N;
428 const ST_L = 4 * N;
429 const ST_B = 6 * N;
430 const ST_F = 7 * N;
431 const ST_FS = 8 * N; // three arrays: spring, bottom and top support forces
432 const ST_AUX = 11 * N; // c[2][3], dc[2][3], cacc[3], reserve
433 const ST_E = ST_AUX + 32; // the strains, ng of them
434 const stLen = ST_E + ng;
435 const AUX_C = 0; // c slot s at AUX_C + 3 s
436 const AUX_DC = 6;
437 const AUX_CACC = 12;
439 // ---- output buffer
440 const OUT_WP = 0;
441 const OUT_WXX = na;
442 const OUT_WYY = 2 * na;
443 const OUT_WXY = 3 * na;
444 const OUT_C = 4 * na;
445 const OUT_FS = OUT_C + nr * na;
446 const OUT_PB = OUT_FS + nsup * na;
447 const OUT_CF = OUT_PB + nbins * na;
448 const outLen = OUT_CF + nr * nf;
450 const s0dt = su.sigma0 * su.dt;
451 const c: Record<string, number> = {
452 INV_DX2: 1 / (su.dx * su.dx),
453 INV_DY2: 1 / (su.dy * su.dy),
454 INV_DXDY4: 1 / (4 * su.dx * su.dy),
455 INV_DT: 1 / su.dt,
456 DT2: su.dt2,
457 CSIG: s0dt / (1 + s0dt),
458 INV_DT2: 1 / su.dt2,
459 BLOWUP: su.blowup,
460 RMAX0: rmax[0],
461 RMAX1: rmax[1],
462 RMAX2: rmax[2],
463 };
464 const ip = su.ip - 1;
465 const jp = su.jp - 1;
466 if (ip < 1 || ip > ny - 2 || jp < 1 || jp > nx - 2)
467 throw new Error("probe is not an interior node");
469 const header = `
470const NX = ${nx}u; const NY = ${ny}u; const N = ${N}u;
471const NSTEPS = ${nsteps}u; const NA = ${na}u;
472const NBINS = ${nbins}u; const NR = ${nr}u; const NSUP = ${nsup}u; const NF = ${nf}u; const NG = ${ng}u;
473const IPROBE = ${ip + jp * ny}u;
474${Object.entries(c)
475 .map(([k, v]) => `const ${k}: f32 = ${fmt(v)};`)
476 .join("\n")}
477${Object.entries(ops.off)
478 .map(([k, v]) => `const OP_${k} = ${v}u;`)
479 .join("\n")}
480${Object.entries(idx.off)
481 .map(([k, v]) => `const IX_${k} = ${v}u;`)
482 .join("\n")}
483const ST_PHI = ${ST_PHI}u; const ST_D = ${ST_D}u; const ST_L = ${ST_L}u; const ST_B = ${ST_B}u;
484const ST_F = ${ST_F}u; const ST_FS = ${ST_FS}u; const ST_AUX = ${ST_AUX}u; const ST_E = ${ST_E}u;
485const AUX_C = ${AUX_C}u; const AUX_DC = ${AUX_DC}u; const AUX_CACC = ${AUX_CACC}u;
486const OUT_WP = ${OUT_WP}u; const OUT_WXX = ${OUT_WXX}u; const OUT_WYY = ${OUT_WYY}u; const OUT_WXY = ${OUT_WXY}u;
487const OUT_C = ${OUT_C}u; const OUT_FS = ${OUT_FS}u; const OUT_PB = ${OUT_PB}u; const OUT_CF = ${OUT_CF}u;
489@group(0) @binding(0) var<storage, read> ops: array<f32>;
490@group(0) @binding(1) var<storage, read> idx: array<u32>;
491@group(0) @binding(2) var<storage, read_write> st: array<f32>;
492@group(0) @binding(3) var<storage, read_write> frames: array<f32>;
493@group(0) @binding(4) var<storage, read_write> outb: array<f32>;
494@group(0) @binding(5) var<storage, read_write> ctl: array<u32>;
495@group(0) @binding(6) var<uniform> uni: vec4<u32>;
497// The blow-up flag holds the 1-based step at which the field passed the limit;
498// the steps after it do nothing, the step itself still records, as the .m does.
499fn stopped(k: u32) -> bool {
500 let d = ctl[1];
501 return d != 0u && d <= k;
502}
503// Slot of the current field (phi, d, Lphi, c, dc) at step k and of the one being written
504fn cur(k: u32) -> u32 { return (k & 1u) * N; }
505fn nxt(k: u32) -> u32 { return ((k + 1u) & 1u) * N; }
506fn curc(k: u32) -> u32 { return (k & 1u) * 3u; }
507fn nxtc(k: u32) -> u32 { return ((k + 1u) & 1u) * 3u; }
508// Rigid coordinate j of the current step, zero without rigid coordinates
509fn cget(k: u32, j: u32) -> f32 { if (NR == 0u) { return 0.0; } return st[ST_AUX + AUX_C + curc(k) + j]; }
510fn dcget(k: u32, j: u32) -> f32 { if (NR == 0u) { return 0.0; } return st[ST_AUX + AUX_DC + curc(k) + j]; }
511fn rget(i: u32, j: u32) -> f32 { if (NR == 0u) { return 0.0; } return ops[OP_R + i + j * N]; }
512// Total field and its backward-difference velocity at node i, step k
513fn wtot(k: u32, i: u32) -> f32 {
514 var w = st[ST_PHI + cur(k) + i];
515 if (NR > 0u) { w = w + rget(i, 0u) * cget(k, 0u) + rget(i, 1u) * cget(k, 1u) + rget(i, 2u) * cget(k, 2u); }
516 return w;
517}
518fn vtot(k: u32, i: u32) -> f32 {
519 var d = st[ST_D + cur(k) + i];
520 if (NR > 0u) { d = d + rget(i, 0u) * dcget(k, 0u) + rget(i, 1u) * dcget(k, 1u) + rget(i, 2u) * dcget(k, 2u); }
521 return d * INV_DT;
522}
523`;
525 // Every difference stencil of the energy method applied to phi; the stencils
526 // are integers, so a strain is a sum of nearby field values with small
527 // integer coefficients, nearly exact in f32.
528 const strain = `${header}
529@compute @workgroup_size(${WG})
530fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
531 let k = ctl[0] + uni.x;
532 if (k >= NSTEPS) { return; }
533 let q = gid.x;
534 if (q >= NG) { return; }
535 if (stopped(k)) { return; }
536 let pS = ST_PHI + cur(k);
537 var e: f32 = 0.0;
538 for (var p = idx[IX_GPTR + q]; p < idx[IX_GPTR + q + 1u]; p = p + 1u) {
539 e = e + ops[OP_GV + p] * st[pS + idx[IX_GCOL + p]];
540 }
541 st[ST_E + q] = e;
542}
543`;
545 // G' applied to the weighted strains: bending, prestress and Laplacian at once
546 const force = `${header}
547@compute @workgroup_size(${WG})
548fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
549 let k = ctl[0] + uni.x;
550 if (k >= NSTEPS) { return; }
551 let i = gid.x;
552 if (i >= N) { return; }
553 if (stopped(k)) { return; }
554 var yb: f32 = 0.0; var yg: f32 = 0.0; var yl: f32 = 0.0;
555 for (var p = idx[IX_GTPTR + i]; p < idx[IX_GTPTR + i + 1u]; p = p + 1u) {
556 let q = idx[IX_GTCOL + p];
557 let g = ops[OP_GTV + p];
558 var sb: f32 = 0.0; var sg: f32 = 0.0; var sl: f32 = 0.0;
559 for (var r = idx[IX_WPTR + q]; r < idx[IX_WPTR + q + 1u]; r = r + 1u) {
560 let e = st[ST_E + idx[IX_WCOL + r]];
561 sb = sb + ops[OP_WBV + r] * e;
562 sg = sg + ops[OP_WGV + r] * e;
563 sl = sl + ops[OP_WLV + r] * e;
564 }
565 yb = yb + g * sb; yg = yg + g * sg; yl = yl + g * sl;
566 }
567 let phi = st[ST_PHI + cur(k) + i];
568 st[ST_B + i] = yb + (yg + ops[OP_KSD + i] * phi);
569 st[ST_L + cur(k) + i] = yl * ops[OP_INVA + i];
570}
571`;
573 // Nodal forces, the reductions of the rigid-body equations and the supports,
574 // the rigid update, and this step's series, all in one workgroup.
575 const dots = `${header}
576var<workgroup> sh0: array<f32, ${RWG}>;
577var<workgroup> sh1: array<f32, ${RWG}>;
578var<workgroup> sh2: array<f32, ${RWG}>;
579var<workgroup> sh3: array<f32, ${RWG}>;
580var<workgroup> sh4: array<f32, ${RWG}>;
581var<workgroup> sh5: array<f32, ${RWG}>;
582var<workgroup> shs: array<f32, ${RWG}>;
583fn reduce6(lid: u32) {
584 for (var s = ${RWG / 2}u; s > 0u; s = s >> 1u) {
585 if (lid < s) {
586 sh0[lid] = sh0[lid] + sh0[lid + s]; sh1[lid] = sh1[lid] + sh1[lid + s]; sh2[lid] = sh2[lid] + sh2[lid + s];
587 sh3[lid] = sh3[lid] + sh3[lid + s]; sh4[lid] = sh4[lid] + sh4[lid + s]; sh5[lid] = sh5[lid] + sh5[lid + s];
588 }
589 workgroupBarrier();
590 }
591}
592@compute @workgroup_size(${RWG})
593fn main(@builtin(local_invocation_index) lid: u32) {
594 let k = ctl[0] + uni.x;
595 let stop = stopped(k) || k >= NSTEPS;
596 let pS = ST_PHI + cur(k);
597 let lS = ST_L + cur(k); let lmS = ST_L + nxt(k);
598 let Fk = ops[OP_F + min(k, NSTEPS - 1u)];
599 var dc0: f32 = dcget(k, 0u); var dc1: f32 = dcget(k, 1u); var dc2: f32 = dcget(k, 2u);
600 // nodal forces about rest, and the partial sums R' f and (AR)' phi
601 var a0: f32 = 0.0; var a1: f32 = 0.0; var a2: f32 = 0.0;
602 var g0: f32 = 0.0; var g1: f32 = 0.0; var g2: f32 = 0.0;
603 for (var i = lid; i < N; i = i + ${RWG}u) {
604 let w = wtot(k, i);
605 let v = vtot(k, i);
606 // the sigma1 loss on the total velocity: 2 sigma1 M Lm (w - wm) / dt
607 var dlw = st[lS + i] - st[lmS + i];
608 if (NR > 0u) { dlw = dlw + ops[OP_LR + i] * dc0 + ops[OP_LR + N + i] * dc1 + ops[OP_LR + 2u * N + i] * dc2; }
609 var f = Fk * ops[OP_SA + i] + ops[OP_S1F + i] * dlw;
610 var fsS = ops[OP_FSTATS + i] - ops[OP_KS + i] * w;
611 var fsB = ops[OP_FSTATB + i];
612 var fsT = ops[OP_FSTATT + i];
613 let kb = ops[OP_KB + i];
614 if (kb > 0.0) {
615 let gb = ops[OP_GB + i];
616 let comp = gb - w;
617 var fb = kb * max(comp, 0.0) - kb * max(gb, 0.0);
618 if (comp > 0.0) { fb = fb - ops[OP_CB + i] * v; }
619 f = f + fb; fsB = fsB + fb;
620 }
621 let kt = ops[OP_KT + i];
622 if (kt > 0.0) {
623 let gt = ops[OP_GT + i];
624 let comp = gt + w;
625 var ft = -kt * max(comp, 0.0) + kt * max(gt, 0.0);
626 if (comp > 0.0) { ft = ft - ops[OP_CT + i] * v; }
627 f = f + ft; fsT = fsT + ft;
628 }
629 let cs = ops[OP_CS + i];
630 if (cs > 0.0) { let fl = -cs * v; f = f + fl; fsS = fsS + fl; }
631 st[ST_F + i] = f;
632 st[ST_FS + i] = fsS;
633 st[ST_FS + N + i] = fsB;
634 st[ST_FS + 2u * N + i] = fsT;
635 if (NR > 0u) {
636 let phi = st[pS + i];
637 a0 = a0 + rget(i, 0u) * f; a1 = a1 + rget(i, 1u) * f; a2 = a2 + rget(i, 2u) * f;
638 g0 = g0 + ops[OP_AR + i] * phi; g1 = g1 + ops[OP_AR + N + i] * phi; g2 = g2 + ops[OP_AR + 2u * N + i] * phi;
639 }
640 }
641 sh0[lid] = a0; sh1[lid] = a1; sh2[lid] = a2; sh3[lid] = g0; sh4[lid] = g1; sh5[lid] = g2;
642 workgroupBarrier();
643 reduce6(lid);
644 // rigid-body update and the rigid series, on one thread
645 if (lid == 0u && NR > 0u) {
646 let cS = ST_AUX + AUX_C + curc(k); let cnS = ST_AUX + AUX_C + nxtc(k);
647 let dS = ST_AUX + AUX_DC + curc(k); let dnS = ST_AUX + AUX_DC + nxtc(k);
648 var rf = array<f32, 3>(sh0[0], sh1[0], sh2[0]);
649 var g = array<f32, 3>(sh3[0], sh4[0], sh5[0]);
650 var cc = array<f32, 3>(st[cS], st[cS + 1u], st[cS + 2u]);
651 var dc = array<f32, 3>(st[dS], st[dS + 1u], st[dS + 2u]);
652 var tmp = array<f32, 3>(0.0, 0.0, 0.0);
653 for (var p = 0u; p < 3u; p = p + 1u) {
654 var s = rf[p] - g[p];
655 for (var q = 0u; q < 3u; q = q + 1u) { s = s - ops[OP_KRR + p + 3u * q] * cc[q]; }
656 tmp[p] = s;
657 }
658 if (!stop) {
659 let ia = idx[IX_IA + k];
660 for (var p = 0u; p < 3u; p = p + 1u) {
661 var acc: f32 = 0.0;
662 for (var q = 0u; q < 3u; q = q + 1u) { acc = acc + ops[OP_GINV + p + 3u * q] * tmp[q]; }
663 let rest = DT2 * acc;
664 let dn = (dc[p] + rest) - CSIG * (2.0 * dc[p] + rest);
665 st[ST_AUX + AUX_CACC + p] = acc;
666 st[dnS + p] = dn;
667 st[cnS + p] = cc[p] + dn;
668 outb[OUT_C + p * NA + ia] = outb[OUT_C + p * NA + ia] + cc[p] * ops[OP_WSAMP + ia];
669 }
670 }
671 }
672 // the probe series of this step
673 if (lid == 0u && !stop) {
674 let ia = idx[IX_IA + k];
675 let cw = st[pS + IPROBE];
676 let e = st[pS + IPROBE + NY]; let w = st[pS + IPROBE - NY];
677 let n = st[pS + IPROBE + 1u]; let so = st[pS + IPROBE - 1u];
678 let ne = st[pS + IPROBE + NY + 1u]; let nw = st[pS + IPROBE - NY + 1u];
679 let se = st[pS + IPROBE + NY - 1u]; let sw = st[pS + IPROBE - NY - 1u];
680 outb[OUT_WP + ia] = outb[OUT_WP + ia] + wtot(k, IPROBE) * ops[OP_WSAMP + ia];
681 outb[OUT_WXX + ia] = outb[OUT_WXX + ia] + (e - 2.0 * cw + w) * INV_DX2 * ops[OP_WSAMP + ia];
682 outb[OUT_WYY + ia] = outb[OUT_WYY + ia] + (n - 2.0 * cw + so) * INV_DY2 * ops[OP_WSAMP + ia];
683 outb[OUT_WXY + ia] = outb[OUT_WXY + ia] + (ne - se - nw + sw) * INV_DXDY4 * ops[OP_WSAMP + ia];
684 }
685 // support forces, one reduction per support
686 for (var q = 0u; q < NSUP; q = q + 1u) {
687 workgroupBarrier();
688 let p0 = idx[IX_SUPPTR + q]; let p1 = idx[IX_SUPPTR + q + 1u];
689 let base = ST_FS + idx[IX_SUPKIND + q] * N;
690 var a: f32 = 0.0;
691 for (var p = p0 + lid; p < p1; p = p + ${RWG}u) { a = a + st[base + idx[IX_SUPNODES + p]]; }
692 shs[lid] = a;
693 workgroupBarrier();
694 for (var s = ${RWG / 2}u; s > 0u; s = s >> 1u) {
695 if (lid < s) { shs[lid] = shs[lid] + shs[lid + s]; }
696 workgroupBarrier();
697 }
698 if (lid == 0u && !stop) {
699 let ia = idx[IX_IA + k];
700 outb[OUT_FS + q * NA + ia] = outb[OUT_FS + q * NA + ia] + shs[0] * ops[OP_WSAMP + ia];
701 }
702 }
703}
704`;
706 const update = `${header}
707@compute @workgroup_size(${WG})
708fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
709 let k = ctl[0] + uni.x;
710 if (k >= NSTEPS) { return; }
711 let i = gid.x;
712 if (i >= N) { return; }
713 if (stopped(k)) { return; }
714 let pS = ST_PHI + cur(k); let pnS = ST_PHI + nxt(k);
715 let dS = ST_D + cur(k); let dnS = ST_D + nxt(k);
716 var fi = st[ST_F + i] - st[ST_B + i];
717 if (NR > 0u) {
718 // the force of the rigid displacement on the plate, AR c, with c of the current step
719 fi = fi - ops[OP_AR + i] * cget(k, 0u) - ops[OP_AR + N + i] * cget(k, 1u) - ops[OP_AR + 2u * N + i] * cget(k, 2u);
720 }
721 var acc = fi * ops[OP_MINV + i];
722 if (NR > 0u) {
723 acc = acc - rget(i, 0u) * st[ST_AUX + AUX_CACC] - rget(i, 1u) * st[ST_AUX + AUX_CACC + 1u] - rget(i, 2u) * st[ST_AUX + AUX_CACC + 2u];
724 }
725 let rest = DT2 * acc;
726 let d = st[dS + i];
727 var dn = (d + rest) - CSIG * (2.0 * d + rest);
728 var pn = st[pS + i] + dn;
729 let kp = ops[OP_KEEP + i];
730 pn = pn * kp; dn = dn * kp;
731 st[pnS + i] = pn;
732 st[dnS + i] = dn;
733 let f = idx[IX_FRAME + k];
734 if (f != ${NONE}u) {
735 frames[f * N + i] = pn;
736 if (NR > 0u && i < 3u) { outb[OUT_CF + i * NF + f] = st[ST_AUX + AUX_C + nxtc(k) + i]; }
737 }
738}
739`;
741 // Workgroups 0..NBINS-1 sum the weighted acceleration of their bin's entries
742 // into the audio sample of this step; workgroup NBINS takes the extent of the
743 // new field for the blow-up test.
744 const reduce = `${header}
745var<workgroup> sh: array<f32, ${RWG}>;
746@compute @workgroup_size(${RWG})
747fn main(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wgid: vec3<u32>) {
748 let k = ctl[0] + uni.x;
749 let stop = stopped(k) || k >= NSTEPS;
750 let pnS = ST_PHI + nxt(k);
751 let dS = ST_D + cur(k); let dnS = ST_D + nxt(k);
752 let wg = wgid.x;
753 if (wg < NBINS) {
754 var dd = array<f32, 3>(0.0, 0.0, 0.0);
755 if (NR > 0u) {
756 for (var p = 0u; p < 3u; p = p + 1u) { dd[p] = st[ST_AUX + AUX_DC + nxtc(k) + p] - st[ST_AUX + AUX_DC + curc(k) + p]; }
757 }
758 let p0 = idx[IX_BINPTR + wg]; let p1 = idx[IX_BINPTR + wg + 1u];
759 var a: f32 = 0.0;
760 for (var p = p0 + lid; p < p1; p = p + ${RWG}u) {
761 let i = idx[IX_BINNODE + p];
762 var accw = st[dnS + i] - st[dS + i];
763 if (NR > 0u) { accw = accw + rget(i, 0u) * dd[0] + rget(i, 1u) * dd[1] + rget(i, 2u) * dd[2]; }
764 a = a + ops[OP_WGT + p] * (accw * INV_DT2);
765 }
766 sh[lid] = a;
767 workgroupBarrier();
768 for (var s = ${RWG / 2}u; s > 0u; s = s >> 1u) {
769 if (lid < s) { sh[lid] = sh[lid] + sh[lid + s]; }
770 workgroupBarrier();
771 }
772 if (lid == 0u && !stop) {
773 let ia = idx[IX_IA + k];
774 let o = OUT_PB + wg + ia * NBINS;
775 outb[o] = outb[o] + sh[0] * ops[OP_WSAMP + ia];
776 }
777 } else {
778 var m: f32 = 0.0;
779 for (var i = lid; i < N; i = i + ${RWG}u) { m = max(m, abs(st[pnS + i])); }
780 sh[lid] = m;
781 workgroupBarrier();
782 for (var s = ${RWG / 2}u; s > 0u; s = s >> 1u) {
783 if (lid < s) { sh[lid] = max(sh[lid], sh[lid + s]); }
784 workgroupBarrier();
785 }
786 if (lid == 0u && !stop) {
787 var peak = sh[0];
788 if (NR > 0u) {
789 let cn = ST_AUX + AUX_C + nxtc(k);
790 peak = peak + abs(st[cn]) * RMAX0 + abs(st[cn + 1u]) * RMAX1 + abs(st[cn + 2u]) * RMAX2;
791 }
792 if (!(peak < BLOWUP)) { ctl[1] = k + 1u; }
793 }
794 }
795}
796`;
798 const sources = [strain, force, dots, update, reduce];
799 const groups = [Math.ceil(ng / WG), Math.ceil(N / WG), 1, Math.ceil(N / WG), nbins + 1];
801 // ---- initial state, in f64 as the .m has it
802 const st0 = new Float32Array(stLen);
803 const e0 = spmv(su.Gs, su.phi0, ng);
804 const wl0 = spmv(su.WL, e0, ng);
805 const lphi0 = spmvT(su.Gs, wl0, N);
806 for (let i = 0; i < N; i++) lphi0[i] *= su.inva[i];
807 for (let i = 0; i < N; i++) {
808 st0[ST_PHI + i] = su.phi0[i];
809 st0[ST_L + i] = lphi0[i];
810 st0[ST_L + N + i] = lphi0[i];
811 }
812 for (let j = 0; j < nr; j++) st0[ST_AUX + AUX_C + j] = su.c0[j];
814 return {
815 N,
816 nf,
817 nr,
818 nsup,
819 nbins,
820 na,
821 outLen,
822 stLen,
823 st0,
824 ops: ops.data,
825 idx: idx.data,
826 sources,
827 groups,
828 out: { OUT_WP, OUT_WXX, OUT_WYY, OUT_WXY, OUT_C, OUT_FS, OUT_PB, OUT_CF },
829 };
830}
832export async function runFdLoopGpu(
833 device: GPUDevice,
834 su: FdLoopSetup,
835 onProgress?: (stepsDone: number, nsteps: number) => void,
836): Promise<FdLoopRecord> {
837 const nsteps = su.nsteps | 0;
838 const b = buildShaders(su);
840 const bgl = device.createBindGroupLayout({
841 entries: [
842 { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
843 { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
844 { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
845 { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
846 { binding: 4, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
847 { binding: 5, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
848 {
849 binding: 6,
850 visibility: GPUShaderStage.COMPUTE,
851 buffer: { type: "uniform", hasDynamicOffset: true, minBindingSize: 16 },
852 },
853 ],
854 });
855 const layout = device.createPipelineLayout({ bindGroupLayouts: [bgl] });
856 const pipelines: GPUComputePipeline[] = [];
857 for (const code of b.sources) {
858 const module = device.createShaderModule({ code });
859 const info = await module.getCompilationInfo();
860 const errs = info.messages.filter((m) => m.type === "error");
861 if (errs.length)
862 throw new Error(
863 "WGSL: " + errs.map((m) => `${m.lineNum}:${m.linePos} ${m.message}`).join("; "),
864 );
865 pipelines.push(
866 device.createComputePipeline({ layout, compute: { module, entryPoint: "main" } }),
867 );
868 }
870 const mk = (size: number, usage: number) =>
871 device.createBuffer({ size: Math.max(16, size), usage });
872 const opsBuf = mk(b.ops.byteLength, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
873 device.queue.writeBuffer(opsBuf, 0, b.ops);
874 const idxBuf = mk(b.idx.byteLength, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
875 device.queue.writeBuffer(idxBuf, 0, b.idx);
876 const stBuf = mk(b.stLen * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
877 device.queue.writeBuffer(stBuf, 0, b.st0);
878 const framesBuf = mk(b.N * b.nf * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
879 const outBuf = mk(b.outLen * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
880 const ctlBuf = mk(16, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST);
881 const uniBuf = mk(CHUNK * 256, GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
882 {
883 const uni = new Uint32Array(CHUNK * 64);
884 for (let s = 0; s < CHUNK; s++) uni[s * 64] = s;
885 device.queue.writeBuffer(uniBuf, 0, uni);
886 }
887 const ctlStage = [0, 1].map(() => mk(16, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST));
888 const readFrames = mk(b.N * b.nf * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
889 const readOut = mk(b.outLen * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
890 const all = [
891 opsBuf,
892 idxBuf,
893 stBuf,
894 framesBuf,
895 outBuf,
896 ctlBuf,
897 uniBuf,
898 ...ctlStage,
899 readFrames,
900 readOut,
901 ];
903 const bindGroup = device.createBindGroup({
904 layout: bgl,
905 entries: [
906 { binding: 0, resource: { buffer: opsBuf } },
907 { binding: 1, resource: { buffer: idxBuf } },
908 { binding: 2, resource: { buffer: stBuf } },
909 { binding: 3, resource: { buffer: framesBuf } },
910 { binding: 4, resource: { buffer: outBuf } },
911 { binding: 5, resource: { buffer: ctlBuf } },
912 { binding: 6, resource: { buffer: uniBuf, offset: 0, size: 16 } },
913 ],
914 });
916 try {
917 const t0 = performance.now();
918 let divergedStep = 0;
919 const pending: { done: Promise<undefined>; stage: GPUBuffer; kend: number }[] = [];
920 const settle = async () => {
921 const p = pending.shift()!;
922 await p.done;
923 await p.stage.mapAsync(GPUMapMode.READ);
924 const flag = new Uint32Array(p.stage.getMappedRange().slice(0))[1];
925 p.stage.unmap();
926 onProgress?.(p.kend, nsteps);
927 if (flag !== 0) divergedStep = flag;
928 };
929 let kbase = 0;
930 let c = 0;
931 while (kbase < nsteps && divergedStep === 0) {
932 const n = Math.min(CHUNK, nsteps - kbase);
933 device.queue.writeBuffer(ctlBuf, 0, Uint32Array.of(kbase));
934 const enc = device.createCommandEncoder();
935 const pass = enc.beginComputePass();
936 for (let s = 0; s < n; s++) {
937 pass.setBindGroup(0, bindGroup, [s * 256]);
938 for (let q = 0; q < pipelines.length; q++) {
939 pass.setPipeline(pipelines[q]);
940 pass.dispatchWorkgroups(b.groups[q]);
941 }
942 }
943 pass.end();
944 const stage = ctlStage[c % 2];
945 enc.copyBufferToBuffer(ctlBuf, 0, stage, 0, 16);
946 device.queue.submit([enc.finish()]);
947 pending.push({ done: device.queue.onSubmittedWorkDone(), stage, kend: kbase + n });
948 kbase += n;
949 c++;
950 if (pending.length >= 2) await settle();
951 }
952 while (pending.length) await settle();
954 const enc = device.createCommandEncoder();
955 enc.copyBufferToBuffer(framesBuf, 0, readFrames, 0, Math.max(16, b.N * b.nf * 4));
956 enc.copyBufferToBuffer(outBuf, 0, readOut, 0, Math.max(16, b.outLen * 4));
957 device.queue.submit([enc.finish()]);
958 await readFrames.mapAsync(GPUMapMode.READ);
959 await readOut.mapAsync(GPUMapMode.READ);
960 const Wf = new Float32Array(readFrames.getMappedRange().slice(0)).subarray(0, b.N * b.nf);
961 const out = new Float32Array(readOut.getMappedRange().slice(0));
962 readFrames.unmap();
963 readOut.unmap();
964 const gpuMs = performance.now() - t0;
965 const o = b.out;
966 const na = b.na;
967 return {
968 w_p: out.slice(o.OUT_WP, o.OUT_WP + na),
969 wxx_p: out.slice(o.OUT_WXX, o.OUT_WXX + na),
970 wyy_p: out.slice(o.OUT_WYY, o.OUT_WYY + na),
971 wxy_p: out.slice(o.OUT_WXY, o.OUT_WXY + na),
972 C: out.slice(o.OUT_C, o.OUT_C + b.nr * na),
973 Fsup: out.slice(o.OUT_FS, o.OUT_FS + b.nsup * na),
974 Wf: Float32Array.from(Wf),
975 Cf: out.slice(o.OUT_CF, o.OUT_CF + b.nr * b.nf),
976 Pbins: out.slice(o.OUT_PB, o.OUT_PB + b.nbins * na),
977 divergedStep,
978 gpuMs,
979 };
980 } finally {
981 for (const buf of all) buf.destroy();
982 }
983}