81f3b30Elliptical plates, a WebGPU time loop for the explicit method, and WAV exportJeremy Magland 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 three compute dispatches for a rectangle with supported or
9 * clamped edges, and three or five when the setup carries sparse operators
10 * (free edges, or an ellipse of any edge condition); the dispatch boundaries
11 * are the global synchronization points the stencil needs:
12 * op Lw = Laplacian of w at every node (ghost layer by reflection), or,
13 * with sparse operators, Lw = Lm w / mvec and B = Km w / mvec
14 * update the new field, wn = c_new (2 w - c_old wm - k2dt2 B + s1dt2 (Lw - Lw_m)
15 * + fdt2 F(n) s); edges zeroed for the dense rectangle, the nodes
16 * that are not unknowns (S.keep = 0) zeroed for the sparse form
17 * dots (nr > 0) the rigid-body components Qb' wn
18 * project (nr > 0) wn - Rb (Qb' wn)
19 * reduce acceleration summed per listener delay bin into the audio bins,
20 * the probe displacement and curvatures, max |wn| for the blow-up test
21 *
22 * Steps are recorded in chunks of a few thousand into one command buffer each,
23 * with the step index a chunk base (ctl[0], written by the host once per chunk)
24 * plus a within-chunk offset from a dynamic-offset uniform; every recorded
25 * dispatch is identical. Two chunks are in flight at a time and the blow-up
26 * flag is read back after each, so a diverging run stops a chunk late rather
27 * than at the end.
28 *
29 * Precision. The GPU works in single precision. The state is carried as
30 * (w, d = w - wm) rather than (w, wm): the update is algebraically the .m's,
31 * dn = (d + rest) - cs (2 d + rest), wn = w + dn, cs = sigma0 dt / (1 + sigma0 dt),
32 * with rest = -k2dt2 B + s1dt2 (Lw - Lw_m) + fdt2 F(n) s, so that no factor of
33 * the form 1 - epsilon is stored in f32 and the rounding of w (about 6e-8 of a
34 * micrometre-scale displacement) does not feed back into the velocity as it
35 * would through 2 w - wm. The acceleration for the audio is (dn - d) / dt^2.
36 */
38export interface Triplets {
39 /** 1-based row and column indices, as MATLAB's find returns them */
40 i: Float64Array;
41 j: Float64Array;
42 v: Float64Array;
43}
45/** The fields of plate_fd_setup's struct that the loop reads; indices are 1-based as in the .m. */
46export interface FdLoopSetup {
47 nx: number;
48 ny: number;
49 nsteps: number;
50 dx: number;
51 dy: number;
52 dt: number;
53 fs: number;
54 bc: string;
55 /** Load footprint, ny x nx column-major */
56 s: Float64Array;
57 /** Force at the nsteps + 1 time levels */
58 F: Float64Array;
59 /** Probe row and column */
60 ip: number;
61 jp: number;
62 /** Steps at which the field is kept, sorted and unique */
63 frameSteps: Float64Array;
64 /** Listener delay bin and weight of every node */
65 bins: Float64Array;
66 wgt: Float64Array;
67 nbins: number;
68 na: number;
69 /** The loop uses the sparse operators (free edges, or an ellipse) */
70 sparse: boolean;
71 /** Sparse form only: 1 at the unknowns and 0 at the nodes held at zero (N), node weights (N) */
72 keep: Float64Array | null;
73 mvec: Float64Array | null;
74 /** Rigid-body motions to remove: basis and projector (N x nr), nr = 0, 2 or 3 */
75 Rb: Float64Array | null;
76 Qb: Float64Array | null;
77 nr: number;
78 Km: Triplets | null;
79 Lm: Triplets | null;
80 sigma0: number;
81 k2dt2: number;
82 s1dt2: number;
83 fdt2: number;
84 blowup: number;
85}
87/** What plate_fd_loop.m returns, in single precision, before conversion for plate_fd_output.m. */
88export interface FdLoopRecord {
89 w_p: Float32Array;
90 wxx_p: Float32Array;
91 wyy_p: Float32Array;
92 wxy_p: Float32Array;
93 /** ny x nx x nf */
94 Wf: Float32Array;
95 /** nbins x na */
96 Pbins: Float32Array;
97 divergedStep: number;
98 gpuMs: number;
99}
101export interface GpuHandle {
102 device: GPUDevice;
103 adapter: string;
104}
106export function describeAdapter(adapter: GPUAdapter): string {
107 const info = (adapter as { info?: Partial<GPUAdapterInfo> }).info;
108 if (!info) return "unknown adapter";
109 const name = [info.vendor, info.architecture].filter(Boolean).join(" ");
110 return name || info.description || "unknown adapter";
111}
113/** The device request the browser worker and the Node test both make. */
114export async function requestGpuDevice(): Promise<GpuHandle> {
115 const gpu = (globalThis as { navigator?: { gpu?: GPU } }).navigator?.gpu;
116 if (!gpu) throw new Error("WebGPU is not available");
117 const adapter = await gpu.requestAdapter();
118 if (!adapter) throw new Error("no WebGPU adapter");
119 const device = await adapter.requestDevice({
120 requiredLimits: {
121 maxStorageBufferBindingSize: adapter.limits.maxStorageBufferBindingSize,
122 maxBufferSize: adapter.limits.maxBufferSize,
123 },
124 });
125 return { device, adapter: describeAdapter(adapter) };
126}
128const WG = 64;
129const RWG = 256;
130const CHUNK = 2048;
131const NONE = 0xffffffff;
133const fmt = (x: number): string => {
134 if (!Number.isFinite(x)) throw new Error("non-finite kernel constant");
135 let s = x.toPrecision(9);
136 if (!/[.e]/.test(s)) s += ".0";
137 return s + "f";
138};
140interface Csr {
141 ptr: Uint32Array;
142 col: Uint32Array;
143 val: Float32Array;
144}
146function toCsr(t: Triplets, n: number): Csr {
147 const count = new Uint32Array(n);
148 for (let p = 0; p < t.i.length; p++) count[t.i[p] - 1]++;
149 const ptr = new Uint32Array(n + 1);
150 for (let r = 0; r < n; r++) ptr[r + 1] = ptr[r] + count[r];
151 const fill = ptr.slice(0, n);
152 const col = new Uint32Array(t.i.length);
153 const val = new Float32Array(t.i.length);
154 for (let p = 0; p < t.i.length; p++) {
155 const q = fill[t.i[p] - 1]++;
156 col[q] = t.j[p] - 1;
157 val[q] = t.v[p];
158 }
159 return { ptr, col, val };
160}
162/** Nodes grouped by listener delay bin, as a CSR-like list. */
163function binLists(bins: Float64Array, nbins: number): { ptr: Uint32Array; nodes: Uint32Array } {
164 const n = bins.length;
165 const count = new Uint32Array(nbins);
166 for (let i = 0; i < n; i++) count[bins[i] - 1]++;
167 const ptr = new Uint32Array(nbins + 1);
168 for (let b = 0; b < nbins; b++) ptr[b + 1] = ptr[b] + count[b];
169 const fill = ptr.slice(0, nbins);
170 const nodes = new Uint32Array(n);
171 for (let i = 0; i < n; i++) nodes[fill[bins[i] - 1]++] = i;
172 return { ptr, nodes };
173}
175function pack<T extends Uint32Array | Float32Array>(
176 parts: [string, ArrayLike<number>][],
177 make: (n: number) => T,
178): { data: T; off: Record<string, number> } {
179 const off: Record<string, number> = {};
180 let cursor = 0;
181 for (const [name, arr] of parts) {
182 off[name] = cursor;
183 cursor += arr.length;
184 }
185 const data = make(Math.max(1, cursor));
186 for (const [name, arr] of parts) data.set(arr, off[name]);
187 return { data, off };
188}
190function buildShaders(su: FdLoopSetup) {
191 const nx = su.nx | 0;
192 const ny = su.ny | 0;
193 const N = nx * ny;
194 const nsteps = su.nsteps | 0;
195 const nbins = su.nbins | 0;
196 const na = su.na | 0;
197 const nf = su.frameSteps.length;
198 const sp = su.sparse;
199 const nr = sp ? su.nr : 0;
200 // With a projection the frames are written after it; otherwise straight from the update
201 const framesInUpdate = nr === 0;
202 if (su.s.length !== N || su.bins.length !== N || su.wgt.length !== N)
203 throw new Error("setup arrays do not match the grid");
204 if (su.F.length < nsteps) throw new Error("force history shorter than the step count");
206 // ---- per-step tables, decided in f64 exactly as the .m decides them
207 const frameOf = new Uint32Array(nsteps).fill(NONE);
208 for (let kf = 0; kf < nf; kf++) {
209 const step = su.frameSteps[kf];
210 if (step >= 1 && step <= nsteps) frameOf[step - 1] = kf;
211 }
212 const iaOf = new Uint32Array(nsteps);
213 for (let k = 0; k < nsteps; k++) {
214 const nstep = k + 1;
215 const ia = Math.floor(nstep * su.dt * su.fs); // ia = floor(nstep * dt * fs) + 1, made 0-based
216 if (ia >= na) throw new Error("audio bin index out of range");
217 iaOf[k] = ia;
218 }
219 const bl = binLists(su.bins, nbins);
221 const f32 = (a: ArrayLike<number>) => Float32Array.from(a);
222 const opsParts: [string, ArrayLike<number>][] = [
223 ["S", f32(su.s)],
224 ["WGT", f32(su.wgt)],
225 ["F", f32(su.F.subarray(0, nsteps))],
226 ];
227 const idxParts: [string, ArrayLike<number>][] = [
228 ["BINPTR", bl.ptr],
229 ["BINNODES", bl.nodes],
230 ["FRAME", frameOf],
231 ["IA", iaOf],
232 ];
233 if (sp) {
234 if (!su.mvec || !su.keep || !su.Km || !su.Lm)
235 throw new Error("the sparse form needs mvec, keep, Km and Lm");
236 if (su.mvec.length !== N || su.keep.length !== N)
237 throw new Error("node arrays do not match the grid");
238 const K = toCsr(su.Km, N);
239 const L = toCsr(su.Lm, N);
240 opsParts.push(["MVEC", f32(su.mvec)], ["KEEP", f32(su.keep)], ["KV", K.val], ["LV", L.val]);
241 idxParts.push(["KPTR", K.ptr], ["KCOL", K.col], ["LPTR", L.ptr], ["LCOL", L.col]);
242 if (nr > 0) {
243 if (!su.Rb || !su.Qb || su.Rb.length !== N * nr || su.Qb.length !== N * nr)
244 throw new Error("rigid-body arrays do not match the grid");
245 opsParts.push(["RB", f32(su.Rb)], ["QB", f32(su.Qb)]);
246 }
247 }
248 const ops = pack(opsParts, (n) => new Float32Array(n));
249 const idx = pack(idxParts, (n) => new Uint32Array(n));
251 // ---- state buffer: two slots each of w, d = w - wm and Lw; B; a few scalars
252 const ST_W = 0;
253 const ST_D = 2 * N;
254 const ST_L = 4 * N;
255 const ST_B = 6 * N;
256 const ST_AUX = 7 * N;
257 const stLen = ST_AUX + 16;
259 const nprobe = nsteps + 1;
260 const OUT_PB = 4 * nprobe;
261 const outLen = OUT_PB + nbins * na;
263 const s0dt = su.sigma0 * su.dt;
264 const c: Record<string, number> = {
265 INV_DX2: 1 / (su.dx * su.dx),
266 INV_DY2: 1 / (su.dy * su.dy),
267 INV_DXDY4: 1 / (4 * su.dx * su.dy),
268 K2DT2: su.k2dt2,
269 S1DT2: su.s1dt2,
270 FDT2: su.fdt2,
271 CS: s0dt / (1 + s0dt),
272 INV_DT2: 1 / (su.dt * su.dt),
273 DTFS: su.dt * su.fs,
274 BLOWUP: su.blowup,
275 SGN: su.bc === "clamped" ? 1 : -1,
276 };
277 const ip = su.ip - 1;
278 const jp = su.jp - 1;
279 if (ip < 1 || ip > ny - 2 || jp < 1 || jp > nx - 2)
280 throw new Error("probe is not an interior node");
282 const header = `
283const NX = ${nx}u; const NY = ${ny}u; const N = ${N}u;
284const NXI: i32 = ${nx}; const NYI: i32 = ${ny};
285const NSTEPS = ${nsteps}u; const NPROBE = ${nprobe}u;
286const NBINS = ${nbins}u; const NA = ${na}u;
287const IPROBE = ${ip + jp * ny}u;
288const NR = ${nr}u;
289${Object.entries(c)
290 .map(([k, v]) => `const ${k}: f32 = ${fmt(v)};`)
291 .join("\n")}
292${Object.entries(ops.off)
293 .map(([k, v]) => `const OP_${k} = ${v}u;`)
294 .join("\n")}
295${Object.entries(idx.off)
296 .map(([k, v]) => `const IX_${k} = ${v}u;`)
297 .join("\n")}
298const ST_W = ${ST_W}u; const ST_D = ${ST_D}u; const ST_L = ${ST_L}u; const ST_B = ${ST_B}u; const ST_AUX = ${ST_AUX}u;
299const OUT_PB = ${OUT_PB}u;
301@group(0) @binding(0) var<storage, read> ops: array<f32>;
302@group(0) @binding(1) var<storage, read> idx: array<u32>;
303@group(0) @binding(2) var<storage, read_write> st: array<f32>;
304@group(0) @binding(3) var<storage, read_write> frames: array<f32>;
305@group(0) @binding(4) var<storage, read_write> outb: array<f32>;
306@group(0) @binding(5) var<storage, read_write> ctl: array<u32>;
307@group(0) @binding(6) var<uniform> uni: vec4<u32>;
309// The blow-up flag holds the 1-based step at which |w| passed the limit; the
310// steps after it do nothing, the step itself still records, as the .m does.
311fn stopped(k: u32) -> bool {
312 let d = ctl[1];
313 return d != 0u && d <= k;
314}
315// Slot of the current field (w, d, Lw) at step k and of the one being written
316fn cur(k: u32) -> u32 { return (k & 1u) * N; }
317fn nxt(k: u32) -> u32 { return ((k + 1u) & 1u) * N; }
318`;
320 // Padded field of plate_pad: the edge nodes are zero and one ghost layer
321 // reflects the first interior value, odd for supported, even for clamped.
322 const opSupported = `${header}
323fn wat(wS: u32, i: i32, j: i32) -> f32 {
324 var ii = i; var jj = j; var sg: f32 = 1.0;
325 if (ii < 0) { ii = 1; sg = sg * SGN; } else if (ii >= NYI) { ii = NYI - 2; sg = sg * SGN; }
326 if (jj < 0) { jj = 1; sg = sg * SGN; } else if (jj >= NXI) { jj = NXI - 2; sg = sg * SGN; }
327 return sg * st[wS + u32(ii) + u32(jj) * NY];
328}
329@compute @workgroup_size(${WG})
330fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
331 let k = ctl[0] + uni.x;
332 if (k >= NSTEPS) { return; }
333 let i = gid.x;
334 if (i >= N) { return; }
335 if (stopped(k)) { return; }
336 let wS = ST_W + cur(k);
337 let r = i32(i % NY); let cc = i32(i / NY);
338 let w0 = st[wS + i];
339 let lw = (wat(wS, r + 1, cc) - 2.0 * w0 + wat(wS, r - 1, cc)) * INV_DY2
340 + (wat(wS, r, cc + 1) - 2.0 * w0 + wat(wS, r, cc - 1)) * INV_DX2;
341 st[ST_L + cur(k) + i] = lw;
342}
343`;
345 // Sparse form: Lw = Lm w / mvec and B = Km w / mvec from the sparse operators
346 const opSparse = `${header}
347@compute @workgroup_size(${WG})
348fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
349 let k = ctl[0] + uni.x;
350 if (k >= NSTEPS) { return; }
351 let i = gid.x;
352 if (i >= N) { return; }
353 if (stopped(k)) { return; }
354 let wS = ST_W + cur(k);
355 var lsum: f32 = 0.0;
356 for (var p = idx[IX_LPTR + i]; p < idx[IX_LPTR + i + 1u]; p = p + 1u) {
357 lsum = lsum + ops[OP_LV + p] * st[wS + idx[IX_LCOL + p]];
358 }
359 var ksum: f32 = 0.0;
360 for (var p = idx[IX_KPTR + i]; p < idx[IX_KPTR + i + 1u]; p = p + 1u) {
361 ksum = ksum + ops[OP_KV + p] * st[wS + idx[IX_KCOL + p]];
362 }
363 let m = ops[OP_MVEC + i];
364 st[ST_L + cur(k) + i] = lsum / m;
365 st[ST_B + i] = ksum / m;
366}
367`;
369 const biharmonic = sp
370 ? `let b = st[ST_B + i];`
371 : `let lS = ST_L + cur(k);
372 let lw0 = st[lS + i];
373 let b = (st[lS + i + 1u] - 2.0 * lw0 + st[lS + i - 1u]) * INV_DY2
374 + (st[lS + i + NY] - 2.0 * lw0 + st[lS + i - NY]) * INV_DX2;`;
375 const update = `${header}
376@compute @workgroup_size(${WG})
377fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
378 let k = ctl[0] + uni.x;
379 if (k >= NSTEPS) { return; }
380 let i = gid.x;
381 if (i >= N) { return; }
382 if (stopped(k)) { return; }
383 let r = i % NY; let cc = i / NY;
384 let wS = ST_W + cur(k); let wnS = ST_W + nxt(k);
385 let dS = ST_D + cur(k); let dnS = ST_D + nxt(k);
386 var wn: f32 = 0.0; var dn: f32 = 0.0;
387 ${sp ? "if (true) {" : "if (r > 0u && r < NY - 1u && cc > 0u && cc < NX - 1u) {"}
388 ${biharmonic}
389 let lw = st[ST_L + cur(k) + i];
390 let lwm = st[ST_L + nxt(k) + i];
391 let rest = -K2DT2 * b + S1DT2 * (lw - lwm) + FDT2 * ops[OP_F + k] * ops[OP_S + i];
392 let d = st[dS + i];
393 dn = (d + rest) - CS * (2.0 * d + rest);
394 wn = st[wS + i] + dn;
395 }
396 ${sp ? "let kp = ops[OP_KEEP + i]; wn = wn * kp; dn = dn * kp;" : ""}
397 st[wnS + i] = wn;
398 st[dnS + i] = dn;
399 ${
400 framesInUpdate
401 ? `let f = idx[IX_FRAME + k];
402 if (f != ${NONE}u) { frames[f * N + i] = wn; }`
403 : ""
404 }
405}
406`;
408 // Rigid-body components Qb' wn (the nodes held at zero already are)
409 const dots = `${header}
410var<workgroup> sh0: array<f32, ${RWG}>;
411var<workgroup> sh1: array<f32, ${RWG}>;
412var<workgroup> sh2: array<f32, ${RWG}>;
413@compute @workgroup_size(${RWG})
414fn main(@builtin(local_invocation_index) lid: u32) {
415 let k = ctl[0] + uni.x;
416 let wnS = ST_W + nxt(k);
417 var a0: f32 = 0.0; var a1: f32 = 0.0; var a2: f32 = 0.0;
418 for (var i = lid; i < N; i = i + ${RWG}u) {
419 let wn = st[wnS + i];
420 a0 = a0 + ops[OP_QB + i] * wn;
421 a1 = a1 + ops[OP_QB + N + i] * wn;
422 if (NR == 3u) { a2 = a2 + ops[OP_QB + 2u * N + i] * wn; }
423 }
424 sh0[lid] = a0; sh1[lid] = a1; sh2[lid] = a2;
425 workgroupBarrier();
426 for (var s = ${RWG / 2}u; s > 0u; s = s >> 1u) {
427 if (lid < s) { sh0[lid] = sh0[lid] + sh0[lid + s]; sh1[lid] = sh1[lid] + sh1[lid + s]; sh2[lid] = sh2[lid] + sh2[lid + s]; }
428 workgroupBarrier();
429 }
430 if (lid == 0u) { st[ST_AUX] = sh0[0]; st[ST_AUX + 1u] = sh1[0]; st[ST_AUX + 2u] = sh2[0]; }
431}
432`;
434 const project = `${header}
435@compute @workgroup_size(${WG})
436fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
437 let k = ctl[0] + uni.x;
438 if (k >= NSTEPS) { return; }
439 let i = gid.x;
440 if (i >= N) { return; }
441 if (stopped(k)) { return; }
442 let wnS = ST_W + nxt(k); let dnS = ST_D + nxt(k);
443 var corr = ops[OP_RB + i] * st[ST_AUX] + ops[OP_RB + N + i] * st[ST_AUX + 1u];
444 if (NR == 3u) { corr = corr + ops[OP_RB + 2u * N + i] * st[ST_AUX + 2u]; }
445 // Rb is zero at the nodes held at zero; the factor makes that exact in f32 as well
446 let kp = ops[OP_KEEP + i];
447 let wn = (st[wnS + i] - corr) * kp; let dn = (st[dnS + i] - corr) * kp;
448 st[wnS + i] = wn; st[dnS + i] = dn;
449 let f = idx[IX_FRAME + k];
450 if (f != ${NONE}u) { frames[f * N + i] = wn; }
451}
452`;
454 // Workgroups 0..NBINS-1 sum the weighted acceleration of their bin's nodes into
455 // the audio sample of this step; workgroup NBINS takes max |wn| for the
456 // blow-up test and records the probe.
457 const reduce = `${header}
458var<workgroup> sh: array<f32, ${RWG}>;
459@compute @workgroup_size(${RWG})
460fn main(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wgid: vec3<u32>) {
461 let k = ctl[0] + uni.x;
462 let stop = stopped(k) || k >= NSTEPS;
463 let wnS = ST_W + nxt(k);
464 let dS = ST_D + cur(k); let dnS = ST_D + nxt(k);
465 let wg = wgid.x;
466 if (wg < NBINS) {
467 let p0 = idx[IX_BINPTR + wg]; let p1 = idx[IX_BINPTR + wg + 1u];
468 var a: f32 = 0.0;
469 for (var p = p0 + lid; p < p1; p = p + ${RWG}u) {
470 let i = idx[IX_BINNODES + p];
471 a = a + ops[OP_WGT + i] * ((st[dnS + i] - st[dS + i]) * INV_DT2);
472 }
473 sh[lid] = a;
474 workgroupBarrier();
475 for (var s = ${RWG / 2}u; s > 0u; s = s >> 1u) {
476 if (lid < s) { sh[lid] = sh[lid] + sh[lid + s]; }
477 workgroupBarrier();
478 }
479 if (lid == 0u && !stop) {
480 let ia = idx[IX_IA + k];
481 let o = OUT_PB + wg + ia * NBINS;
482 outb[o] = outb[o] + sh[0] * DTFS;
483 }
484 } else {
485 var m: f32 = 0.0;
486 for (var i = lid; i < N; i = i + ${RWG}u) { m = max(m, abs(st[wnS + i])); }
487 sh[lid] = m;
488 workgroupBarrier();
489 for (var s = ${RWG / 2}u; s > 0u; s = s >> 1u) {
490 if (lid < s) { sh[lid] = max(sh[lid], sh[lid + s]); }
491 workgroupBarrier();
492 }
493 if (lid == 0u && !stop) {
494 let c = st[wnS + IPROBE];
495 let e = st[wnS + IPROBE + NY]; let w = st[wnS + IPROBE - NY];
496 let n = st[wnS + IPROBE + 1u]; let so = st[wnS + IPROBE - 1u];
497 let ne = st[wnS + IPROBE + NY + 1u]; let nw = st[wnS + IPROBE - NY + 1u];
498 let se = st[wnS + IPROBE + NY - 1u]; let sw = st[wnS + IPROBE - NY - 1u];
499 outb[k + 1u] = c;
500 outb[NPROBE + k + 1u] = (e - 2.0 * c + w) * INV_DX2;
501 outb[2u * NPROBE + k + 1u] = (n - 2.0 * c + so) * INV_DY2;
502 outb[3u * NPROBE + k + 1u] = (ne - se - nw + sw) * INV_DXDY4;
503 if (sh[0] >= BLOWUP) { ctl[1] = k + 1u; }
504 }
505 }
506}
507`;
509 const op = sp ? opSparse : opSupported;
510 const sources = nr > 0 ? [op, update, dots, project, reduce] : [op, update, reduce];
511 const groups =
512 nr > 0
513 ? [Math.ceil(N / WG), Math.ceil(N / WG), 1, Math.ceil(N / WG), nbins + 1]
514 : [Math.ceil(N / WG), Math.ceil(N / WG), nbins + 1];
515 return { N, nf, nprobe, nbins, na, outLen, stLen, ops: ops.data, idx: idx.data, sources, groups };
516}
518export async function runFdLoopGpu(
519 device: GPUDevice,
520 su: FdLoopSetup,
521 onProgress?: (stepsDone: number, nsteps: number) => void,
522): Promise<FdLoopRecord> {
523 const nsteps = su.nsteps | 0;
524 const b = buildShaders(su);
526 const bgl = device.createBindGroupLayout({
527 entries: [
528 { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
529 { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
530 { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
531 { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
532 { binding: 4, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
533 { binding: 5, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
534 {
535 binding: 6,
536 visibility: GPUShaderStage.COMPUTE,
537 buffer: { type: "uniform", hasDynamicOffset: true, minBindingSize: 16 },
538 },
539 ],
540 });
541 const layout = device.createPipelineLayout({ bindGroupLayouts: [bgl] });
542 const pipelines: GPUComputePipeline[] = [];
543 for (const code of b.sources) {
544 const module = device.createShaderModule({ code });
545 const info = await module.getCompilationInfo();
546 const errs = info.messages.filter((m) => m.type === "error");
547 if (errs.length)
548 throw new Error(
549 "WGSL: " + errs.map((m) => `${m.lineNum}:${m.linePos} ${m.message}`).join("; "),
550 );
551 pipelines.push(
552 device.createComputePipeline({ layout, compute: { module, entryPoint: "main" } }),
553 );
554 }
556 const mk = (size: number, usage: number) =>
557 device.createBuffer({ size: Math.max(16, size), usage });
558 const opsBuf = mk(b.ops.byteLength, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
559 device.queue.writeBuffer(opsBuf, 0, b.ops);
560 const idxBuf = mk(b.idx.byteLength, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
561 device.queue.writeBuffer(idxBuf, 0, b.idx);
562 const stBuf = mk(b.stLen * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST);
563 const framesBuf = mk(b.N * b.nf * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
564 const outBuf = mk(b.outLen * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC);
565 const ctlBuf = mk(16, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST);
566 const uniBuf = mk(CHUNK * 256, GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST);
567 {
568 const uni = new Uint32Array(CHUNK * 64);
569 for (let s = 0; s < CHUNK; s++) uni[s * 64] = s;
570 device.queue.writeBuffer(uniBuf, 0, uni);
571 }
572 const ctlStage = [0, 1].map(() => mk(16, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST));
573 const readFrames = mk(b.N * b.nf * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
574 const readOut = mk(b.outLen * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST);
575 const all = [
576 opsBuf,
577 idxBuf,
578 stBuf,
579 framesBuf,
580 outBuf,
581 ctlBuf,
582 uniBuf,
583 ...ctlStage,
584 readFrames,
585 readOut,
586 ];
588 const bindGroup = device.createBindGroup({
589 layout: bgl,
590 entries: [
591 { binding: 0, resource: { buffer: opsBuf } },
592 { binding: 1, resource: { buffer: idxBuf } },
593 { binding: 2, resource: { buffer: stBuf } },
594 { binding: 3, resource: { buffer: framesBuf } },
595 { binding: 4, resource: { buffer: outBuf } },
596 { binding: 5, resource: { buffer: ctlBuf } },
597 { binding: 6, resource: { buffer: uniBuf, offset: 0, size: 16 } },
598 ],
599 });
601 try {
602 const t0 = performance.now();
603 let divergedStep = 0;
604 const pending: { done: Promise<undefined>; stage: GPUBuffer; kend: number }[] = [];
605 const settle = async () => {
606 const p = pending.shift()!;
607 await p.done;
608 await p.stage.mapAsync(GPUMapMode.READ);
609 const flag = new Uint32Array(p.stage.getMappedRange().slice(0))[1];
610 p.stage.unmap();
611 onProgress?.(p.kend, nsteps);
612 if (flag !== 0) divergedStep = flag;
613 };
614 let kbase = 0;
615 let c = 0;
616 while (kbase < nsteps && divergedStep === 0) {
617 const n = Math.min(CHUNK, nsteps - kbase);
618 device.queue.writeBuffer(ctlBuf, 0, Uint32Array.of(kbase));
619 const enc = device.createCommandEncoder();
620 const pass = enc.beginComputePass();
621 for (let s = 0; s < n; s++) {
622 pass.setBindGroup(0, bindGroup, [s * 256]);
623 for (let q = 0; q < pipelines.length; q++) {
624 pass.setPipeline(pipelines[q]);
625 pass.dispatchWorkgroups(b.groups[q]);
626 }
627 }
628 pass.end();
629 const stage = ctlStage[c % 2];
630 enc.copyBufferToBuffer(ctlBuf, 0, stage, 0, 16);
631 device.queue.submit([enc.finish()]);
632 pending.push({ done: device.queue.onSubmittedWorkDone(), stage, kend: kbase + n });
633 kbase += n;
634 c++;
635 if (pending.length >= 2) await settle();
636 }
637 while (pending.length) await settle();
639 const enc = device.createCommandEncoder();
640 enc.copyBufferToBuffer(framesBuf, 0, readFrames, 0, Math.max(16, b.N * b.nf * 4));
641 enc.copyBufferToBuffer(outBuf, 0, readOut, 0, Math.max(16, b.outLen * 4));
642 device.queue.submit([enc.finish()]);
643 await readFrames.mapAsync(GPUMapMode.READ);
644 await readOut.mapAsync(GPUMapMode.READ);
645 const Wf = new Float32Array(readFrames.getMappedRange().slice(0)).subarray(0, b.N * b.nf);
646 const out = new Float32Array(readOut.getMappedRange().slice(0));
647 readFrames.unmap();
648 readOut.unmap();
649 const gpuMs = performance.now() - t0;
650 const np = b.nprobe;
651 return {
652 w_p: out.slice(0, np),
653 wxx_p: out.slice(np, 2 * np),
654 wyy_p: out.slice(2 * np, 3 * np),
655 wxy_p: out.slice(3 * np, 4 * np),
656 Wf: Float32Array.from(Wf),
657 Pbins: out.slice(4 * np, 4 * np + b.nbins * b.na),
658 divergedStep,
659 gpuMs,
660 };
661 } finally {
662 for (const buf of all) buf.destroy();
663 }
664}