/// // WebGPU implementation of the cymbal2 time loop (cymbal2_sim.m), ported from // the prototype in proof-of-concept/cymbal2-webgpu. Everything around the // loop stays in numbl: the caller runs the unmodified MATLAB setup stage // (grid, per-wavenumber inverses, mallet, stroke schedule) and hands the // tensors in through GpuSetup; the acceleration and contact-force histories // come back for the unmodified output stage. // // One time step is eight compute dispatches at full GPU width, the dispatch // boundaries providing global synchronization between stages: // AB hand state machine + Hertz contact + beam update (workgroup 0), and // the shell derivative stencils / Airy right-hand side (all workgroups) // C azimuthal analysis of the Airy right-hand side // D Airy solve, one small radial system per wavenumber // E azimuthal synthesis of the stress function // FG von Karman bracket + forcing (stress-function stencils inline) // H azimuthal analysis of the nonlinear load // I implicit shell update, three radial products per wavenumber // J azimuthal synthesis of the new displacement; pickup acceleration // // Every recorded dispatch is identical: the step index k is a chunk base // (state cell 0, written by the host once per chunk) plus a within-chunk // offset from a dynamic-offset uniform, and all slot rotation is a pure // function of k. Slots rotate every step whether or not the shell is active, // which is exact: before first contact every slot is zero, and after first // contact every step writes. The shell-active flag latches inside AB and is // read by the later stages of the same step one dispatch later; the benign // race inside AB itself can only skip work on a field that is still zero. export interface GpuSetup { ops: { D1: Float32Array; D2: Float32Array; F: Float32Array; Gi: Float32Array; Rd: Float32Array; R2d: Float32Array; AiB: Float32Array; AiC: Float32Array; Ai: Float32Array; Fi: Float32Array; ASiB: Float32Array; ASiC: Float32Array; vM: Float32Array; vR: Float32Array; vK: Float32Array; tipcol: Float32Array; }; tlaunch: number[]; vstr: number[]; nrt: number; Nr: number; nth: number; hr: number; hth: number; ht: number; Eh: number; Rs: number; rhoh: number; fsc: number; nlr0: number; ls: number; js: number; lp: number; jp: number; tipi: number; nfree: number; ps: number; KM: number; alM: number; aH: number; aHr: number; holdSteps: number; missDepth: number; gap: number; K: number; nsteps: number; } export interface GpuSimResult { acc: Float32Array; fhist: Float32Array; peakXi: number; peakDisp: number; contactVel: number[]; strokesPlayed: number; gpuMs: number; } const fmt = (x: number): string => { if (!Number.isFinite(x)) throw new Error('non-finite kernel constant'); let s = x.toPrecision(9); if (!/[.e]/.test(s)) s += '.0'; return s + 'f'; }; interface Layout { N: number; opsData: Float32Array; stLen: number; SC_CVEL: number; } function buildLayoutAndStages(su: GpuSetup, wg: number) { const nrt = su.nrt | 0; const nth = su.nth | 0; const nf = su.nfree | 0; const N = nrt * nth; const K = su.K | 0; const ht = su.ht; // ---- packed operator buffer ---- const parts: [string, Float32Array, number][] = [ ['D1', su.ops.D1, nrt * nrt], ['D2', su.ops.D2, nrt * nrt], ['F', su.ops.F, nth * nth], ['GI', su.ops.Gi, nth * nth], ['RD', su.ops.Rd, nrt], ['R2D', su.ops.R2d, nrt], ['AIB', su.ops.AiB, nrt * nrt * nth], ['AIC', su.ops.AiC, nrt * nrt * nth], ['AI', su.ops.Ai, nrt * nrt * nth], ['FI', su.ops.Fi, nrt * nrt * nth], ['ASIB', su.ops.ASiB, nf * nf], ['ASIC', su.ops.ASiC, nf * nf], ['VM', su.ops.vM, nf], ['VR', su.ops.vR, nf], ['VK', su.ops.vK, nf], ['TIPCOL', su.ops.tipcol, nf], ]; // Launch-step thresholds: smallest 0-based kk with (kk+1)*ht >= tlaunch, // decided in f64 exactly as MATLAB's `k*ht >= tlaunch` decides it. const kgo = su.tlaunch.map((tl) => { let g = Math.max(0, Math.ceil(tl / ht) - 1); while (g > 0 && g * ht >= tl) g -= 1; while ((g + 1) * ht < tl) g += 1; return g; }); parts.push(['KGO', Float32Array.from(kgo), K]); parts.push(['VSTR', Float32Array.from(su.vstr), K]); // Transposed beam matrices so each beam-row thread streams its own row. const tr = (m: Float32Array) => { const t = new Float32Array(nf * nf); for (let r = 0; r < nf; r++) for (let p = 0; p < nf; p++) t[p + r * nf] = m[r + p * nf]; return t; }; parts.push(['ASIBT', tr(su.ops.ASiB), nf * nf]); parts.push(['ASICT', tr(su.ops.ASiC), nf * nf]); const off: Record = {}; let cursor = 0; for (const [name, arr, len] of parts) { if (arr.length !== len) throw new Error(`ops part ${name}: ${arr.length} != ${len}`); off[name] = cursor; cursor += len; } const opsData = new Float32Array(cursor); for (const [name, arr] of parts) opsData.set(arr, off[name]); // ---- state buffer: 16 scalars, cvel, 3 beam slots, 3+3 field slots, scratch ---- const ST_CVEL = 16; const ST_D = ST_CVEL + K; const ST_U = ST_D + 3 * nf; const ST_UH = ST_U + 3 * N; const scratchNames = ['URR', 'V1', 'V2', 'RHSA', 'RHSAF', 'PH3', 'PHI', 'NL', 'NLF']; const A: Record = {}; let scur = ST_UH + 3 * N; for (const n of scratchNames) { A[n] = scur; scur += N; } const stLen = scur; const c: Record = { HT: ht, HT2I: 1 / (ht * ht), GAP: su.gap, KM: su.KM, ALM: su.alM, PS: su.ps, AH: su.aH, AHR: su.aHr, MISSD: su.missDepth, EHRS: su.Eh / su.Rs, EH2: su.Eh / 2, RSI: 1 / su.Rs, SNL: (ht * ht) / su.rhoh, FSC: su.fsc, C1R: 1 / (2 * su.hr), C2R: 1 / (su.hr * su.hr), C1T: 1 / (2 * su.hth), C2T: 1 / (su.hth * su.hth), }; const header = ` const NRT = ${nrt}u; const NTH = ${nth}u; const N = ${N}u; const NF = ${nf}u; const NSTEPS = ${su.nsteps}u; const KS = ${K}u; const NRU = ${su.Nr}u; const LS0 = ${su.ls - 1}u; const JS0 = ${su.js - 1}u; const LP0 = ${su.lp - 1}u; const JP0 = ${su.jp - 1}u; const TIPI = ${su.tipi - 1}u; const NLR0 = ${su.nlr0 - 1}u; const HOLDSTEPS = ${su.holdSteps | 0}; ${Object.entries(c).map(([k, v]) => `const ${k}: f32 = ${fmt(v)};`).join('\n')} ${Object.entries(off).map(([k, v]) => `const OP_${k} = ${v}u;`).join('\n')} const ST_CVEL = ${ST_CVEL}u; const ST_D = ${ST_D}u; const ST_U = ${ST_U}u; const ST_UH = ${ST_UH}u; ${scratchNames.map((n) => `const A_${n} = ${A[n]}u;`).join('\n')} @group(0) @binding(0) var ops: array; @group(0) @binding(1) var st: array; @group(0) @binding(2) var outb: array; @group(0) @binding(3) var ctl: vec4; fn thm(j: u32) -> u32 { return (j + NTH - 1u) % NTH; } fn thp(j: u32) -> u32 { return (j + 1u) % NTH; } `; const AB = `${header} var wgF: f32; var wgHA: f32; var wgHB: f32; var wgHC: f32; var shD: array; var shD1: array; fn u_dr(b: u32, l: u32, j: u32) -> f32 { if (l < 1u || l > NRT - 2u) { return 0.0; } return (st[b + (l + 1u) + j * NRT] - st[b + (l - 1u) + j * NRT]) * C1R; } @compute @workgroup_size(${wg}) fn main(@builtin(global_invocation_id) gid: vec3, @builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wgid: vec3) { // No early-out on k: the host records exactly nsteps steps, and the // workgroupBarrier below requires uniform control flow up to it. let k = u32(st[0u]) + ctl.x; let uS = ST_U + (k % 3u) * N; if (wgid.x == 0u) { if (lid == 0u) { var shellA = u32(st[1u]); var inC = u32(st[2u]); var phase = u32(st[3u]); var strokeIdx = u32(st[4u]); var holdLeft = i32(st[5u]); var vtar = st[6u]; var vH = st[7u]; var uHm = st[8u]; var uH0 = st[9u]; var uHp = st[10u]; var peakXi = st[12u]; var peakDisp = st[13u]; let dCur = ST_D + ((k + 2u) % 3u) * NF; let dOld = ST_D + ((k + 1u) % 3u) * NF; if (shellA != 0u && k > 0u && (k % 512u) == 0u) { var pd: f32 = 0.0; for (var l: u32 = 0u; l < NRT; l = l + 1u) { pd = max(pd, abs(st[uS + l + JS0 * NRT])); } if (pd > peakDisp) { peakDisp = pd; } } let xi = st[uS + LS0 + JS0 * NRT] - st[dCur + TIPI] - GAP; var fM: f32 = 0.0; if (xi > 0.0) { fM = KM * pow(xi, ALM); shellA = 1u; if (inC == 0u && strokeIdx < KS && st[ST_CVEL + strokeIdx] == 0.0) { st[ST_CVEL + strokeIdx] = (st[dCur + TIPI] - st[dOld + TIPI]) / HT; } inC = 1u; if (xi > peakXi) { peakXi = xi; } } else { inC = 0u; } outb[NSTEPS + k] = fM; wgF = fM; wgHA = -(uHp - 2.0 * uH0 + uHm); wgHB = -(HT * 0.5) * (uHp - uHm); wgHC = -HT * HT * (PS * uH0 + ((1.0 - PS) * 0.5) * (uHp + uHm)); if (phase == 0u) { if (strokeIdx < KS && f32(k) >= ops[OP_KGO + strokeIdx]) { phase = 1u; vtar = ops[OP_VSTR + strokeIdx]; } } else if (phase == 1u) { vH = max(vH - AH * HT, -vtar); if (inC != 0u) { phase = 3u; holdLeft = HOLDSTEPS; } else if (vH <= -vtar) { phase = 2u; } } else if (phase == 2u) { if (inC != 0u) { phase = 3u; holdLeft = HOLDSTEPS; } else if (uHp <= -MISSD) { phase = 4u; } } else if (phase == 3u) { holdLeft = holdLeft - 1; if (holdLeft <= 0) { phase = 4u; } } else if (phase == 4u) { vH = min(vH + AHR * HT, 0.5 * vtar); if (vH > 0.0 && uHp >= -vH * vH / (2.0 * AHR)) { phase = 5u; } } else { vH = vH - AHR * HT; if (vH <= 0.0 || uHp >= 0.0) { vH = 0.0; phase = 0u; strokeIdx = strokeIdx + 1u; } } uHm = uH0; uH0 = uHp; uHp = uHp + vH * HT; st[1u] = f32(shellA); st[2u] = f32(inC); st[3u] = f32(phase); st[4u] = f32(strokeIdx); st[5u] = f32(holdLeft); st[6u] = vtar; st[7u] = vH; st[8u] = uHm; st[9u] = uH0; st[10u] = uHp; st[12u] = peakXi; st[13u] = peakDisp; st[14u] = fM; } if (lid < NF) { shD[lid] = st[ST_D + ((k + 2u) % 3u) * NF + lid]; shD1[lid] = st[ST_D + ((k + 1u) % 3u) * NF + lid]; } workgroupBarrier(); if (lid < NF) { let r = lid; var sum: f32 = ops[OP_VM + r] * wgHA + ops[OP_VR + r] * wgHB + ops[OP_VK + r] * wgHC + (HT * HT * wgF) * ops[OP_TIPCOL + r]; for (var p: u32 = 0u; p < NF; p = p + 1u) { sum = sum + ops[OP_ASIBT + p + r * NF] * shD[p] + ops[OP_ASICT + p + r * NF] * shD1[p]; } st[ST_D + (k % 3u) * NF + r] = sum; } } let i = gid.x; if (i >= N) { return; } if (st[1u] == 0.0) { return; } let l = i % NRT; let j = i / NRT; var ur: f32 = 0.0; var urr: f32 = 0.0; if (l >= 1u && l <= NRT - 2u) { let um = st[uS + (l - 1u) + j * NRT]; let u0 = st[uS + l + j * NRT]; let up = st[uS + (l + 1u) + j * NRT]; ur = (up - um) * C1R; urr = (um - 2.0 * u0 + up) * C2R; } let ujm = st[uS + l + thm(j) * NRT]; let uj0 = st[uS + l + j * NRT]; let ujp = st[uS + l + thp(j) * NRT]; let v1 = ops[OP_RD + l] * ur + ops[OP_R2D + l] * ((ujm - 2.0 * uj0 + ujp) * C2T); let urt = (u_dr(uS, l, thp(j)) - u_dr(uS, l, thm(j))) * C1T; let v2 = ops[OP_RD + l] * urt - ops[OP_R2D + l] * ((ujp - ujm) * C1T); st[A_URR + i] = urr; st[A_V1 + i] = v1; st[A_V2 + i] = v2; var rv: f32 = 0.0; if (l >= 3u && l <= NRU - 2u) { rv = EHRS * (urr + v1) - EH2 * 2.0 * (urr * v1 - v2 * v2); } st[A_RHSA + i] = rv; } `; const gate = ` let k = u32(st[0u]) + ctl.x; if (k >= NSTEPS) { return; } let i = gid.x; if (i >= N) { return; } if (st[1u] == 0.0) { return; } let l = i % NRT; let j = i / NRT;`; const transform = (src: string, dst: string, mat: string) => `${header} @compute @workgroup_size(${wg}) fn main(@builtin(global_invocation_id) gid: vec3) { ${gate} var sum: f32 = 0.0; for (var q: u32 = 0u; q < NTH; q = q + 1u) { sum = sum + st[${src} + l + q * NRT] * ops[${mat} + q + j * NTH]; } st[${dst} + i] = sum; } `; const D = `${header} @compute @workgroup_size(${wg}) fn main(@builtin(global_invocation_id) gid: vec3) { ${gate} let base = OP_FI + j * NRT * NRT; var sum: f32 = 0.0; for (var p: u32 = 0u; p < NRT; p = p + 1u) { sum = sum + ops[base + l + p * NRT] * st[A_RHSAF + p + j * NRT]; } st[A_PH3 + i] = sum; } `; const FG = `${header} fn p_dr(l: u32, j: u32) -> f32 { if (l < 1u || l > NRT - 2u) { return 0.0; } return (st[A_PHI + (l + 1u) + j * NRT] - st[A_PHI + (l - 1u) + j * NRT]) * C1R; } @compute @workgroup_size(${wg}) fn main(@builtin(global_invocation_id) gid: vec3) { ${gate} var nl: f32 = 0.0; if (l >= NLR0 && l < NRU) { let pm = st[A_PHI + (l - 1u) + j * NRT]; let p0 = st[A_PHI + l + j * NRT]; let pp = st[A_PHI + (l + 1u) + j * NRT]; let pr = (pp - pm) * C1R; let prr = (pm - 2.0 * p0 + pp) * C2R; let pjm = st[A_PHI + l + thm(j) * NRT]; let pjp = st[A_PHI + l + thp(j) * NRT]; let w1 = ops[OP_RD + l] * pr + ops[OP_R2D + l] * ((pjm - 2.0 * p0 + pjp) * C2T); let prt = (p_dr(l, thp(j)) - p_dr(l, thm(j))) * C1T; let w2 = ops[OP_RD + l] * prt - ops[OP_R2D + l] * ((pjp - pjm) * C1T); nl = prr * st[A_V1 + i] + st[A_URR + i] * w1 - 2.0 * w2 * st[A_V2 + i] - (prr + w1) * RSI; if (l == LS0 && j == JS0) { nl = nl - st[14u] * FSC; } nl = nl * SNL; } st[A_NL + i] = nl; } `; const I = `${header} @compute @workgroup_size(${wg}) fn main(@builtin(global_invocation_id) gid: vec3) { ${gate} let h3S = ST_UH + (k % 3u) * N; let h13S = ST_UH + ((k + 2u) % 3u) * N; let h3nS = ST_UH + ((k + 1u) % 3u) * N; let base = j * NRT * NRT; var sum: f32 = 0.0; for (var p: u32 = 0u; p < NRT; p = p + 1u) { let pj = p + j * NRT; sum = sum + ops[OP_AIB + base + l + p * NRT] * st[h3S + pj] + ops[OP_AIC + base + l + p * NRT] * st[h13S + pj] + ops[OP_AI + base + l + p * NRT] * st[A_NLF + pj]; } st[h3nS + i] = sum; } `; const J = `${header} @compute @workgroup_size(${wg}) fn main(@builtin(global_invocation_id) gid: vec3) { ${gate} let h3nS = ST_UH + ((k + 1u) % 3u) * N; let uS = ST_U + (k % 3u) * N; let u1S = ST_U + ((k + 2u) % 3u) * N; let unS = ST_U + ((k + 1u) % 3u) * N; var sum: f32 = 0.0; for (var q: u32 = 0u; q < NTH; q = q + 1u) { sum = sum + st[h3nS + l + q * NRT] * ops[OP_GI + q + j * NTH]; } st[unS + i] = sum; if (l == LP0 && j == JP0) { outb[k] = (sum - 2.0 * st[uS + i] + st[u1S + i]) * HT2I; } } `; const sources = [ AB, transform('A_RHSA', 'A_RHSAF', 'OP_F'), D, transform('A_PH3', 'A_PHI', 'OP_GI'), FG, transform('A_NL', 'A_NLF', 'OP_F'), I, J, ]; const layout: Layout = { N, opsData, stLen, SC_CVEL: ST_CVEL }; return { layout, sources, nwg: Math.ceil(N / wg) }; } export async function requestGpuDevice(): Promise { const gpu = (navigator as { gpu?: GPU }).gpu; if (!gpu) throw new Error('WebGPU is not available in this browser'); const adapter = await gpu.requestAdapter(); if (!adapter) throw new Error('no WebGPU adapter'); return adapter.requestDevice(); } export async function runGpuSim( device: GPUDevice, su: GpuSetup, onProgress?: (stepsDone: number, nsteps: number) => void, ): Promise { const wg = 64; const chunk = 1024; const nsteps = su.nsteps; const { layout, sources, nwg } = buildLayoutAndStages(su, wg); const bgl = device.createBindGroupLayout({ entries: [ { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } }, { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } }, { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } }, { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform', hasDynamicOffset: true, minBindingSize: 16 } }, ], }); const plLayout = device.createPipelineLayout({ bindGroupLayouts: [bgl] }); const pipelines: GPUComputePipeline[] = []; for (const code of sources) { const module = device.createShaderModule({ code }); const info = await module.getCompilationInfo(); const errs = info.messages.filter((m) => m.type === 'error'); if (errs.length) { throw new Error('WGSL: ' + errs.map((m) => `${m.lineNum}:${m.linePos} ${m.message}`).join('; ')); } pipelines.push(device.createComputePipeline({ layout: plLayout, compute: { module, entryPoint: 'main' } })); } const mkBuf = (size: number, usage: number) => device.createBuffer({ size, usage }); const opsBuf = mkBuf(layout.opsData.byteLength, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST); device.queue.writeBuffer(opsBuf, 0, layout.opsData); const stBuf = mkBuf(layout.stLen * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST); const outBuf = mkBuf(2 * nsteps * 4, GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC); const uniBuf = mkBuf(chunk * 256, GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST); { const uni = new Uint32Array(chunk * 64); for (let s = 0; s < chunk; s++) uni[s * 64] = s; device.queue.writeBuffer(uniBuf, 0, uni); } const readOut = mkBuf(2 * nsteps * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST); const readSt = mkBuf(layout.stLen * 4, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST); const bindGroup = device.createBindGroup({ layout: bgl, entries: [ { binding: 0, resource: { buffer: opsBuf } }, { binding: 1, resource: { buffer: stBuf } }, { binding: 2, resource: { buffer: outBuf } }, { binding: 3, resource: { buffer: uniBuf, offset: 0, size: 16 } }, ], }); const t0 = performance.now(); const pending: Promise[] = []; let kbase = 0; while (kbase < nsteps) { const n = Math.min(chunk, nsteps - kbase); const kb = new Float32Array(new ArrayBuffer(4)); kb[0] = kbase; device.queue.writeBuffer(stBuf, 0, kb); const enc = device.createCommandEncoder(); const pass = enc.beginComputePass(); for (let s = 0; s < n; s++) { pass.setBindGroup(0, bindGroup, [s * 256]); for (const p of pipelines) { pass.setPipeline(p); pass.dispatchWorkgroups(nwg); } } pass.end(); device.queue.submit([enc.finish()]); pending.push(device.queue.onSubmittedWorkDone()); if (pending.length >= 2) await pending.shift(); kbase += n; if (onProgress) onProgress(kbase, nsteps); } while (pending.length) await pending.shift(); const gpuMs = performance.now() - t0; const enc = device.createCommandEncoder(); enc.copyBufferToBuffer(outBuf, 0, readOut, 0, 2 * nsteps * 4); enc.copyBufferToBuffer(stBuf, 0, readSt, 0, layout.stLen * 4); device.queue.submit([enc.finish()]); await readOut.mapAsync(GPUMapMode.READ); await readSt.mapAsync(GPUMapMode.READ); const outData = new Float32Array(readOut.getMappedRange().slice(0)); const state = new Float32Array(readSt.getMappedRange().slice(0)); readOut.unmap(); readSt.unmap(); for (const b of [opsBuf, stBuf, outBuf, uniBuf, readOut, readSt]) b.destroy(); const acc = outData.subarray(0, nsteps); for (let i = 0; i < acc.length; i++) { if (!Number.isFinite(acc[i])) { throw new Error( `cymbal2: simulation diverged at t = ${(i * su.ht).toFixed(4)} s ` + '(reduce strike_velocity, or raise n_theta / oversample)', ); } } return { acc, fhist: outData.subarray(nsteps, 2 * nsteps), peakXi: state[12], peakDisp: state[13], contactVel: Array.from(state.subarray(layout.SC_CVEL, layout.SC_CVEL + su.K)), strokesPlayed: Math.min(state[4] + (state[3] > 0 ? 1 : 0), su.K), gpuMs, }; }