10729f3WebGPU time loop for the cymbal2 shell model, with numbl setup and output stagesJeremy Magland 1// Persistent-kernel backend: ONE dispatch of NWG workgroups runs CHUNK time
2// steps, with a global spin barrier on storage atomics between stages in
3// place of dispatch boundaries. This removes the per-dispatch launch bubble
4// that keeps an idle-clocked GPU at its minimum frequency, and it uses the
5// whole GPU, unlike the one-workgroup megakernel.
6//
7// WebGPU does not guarantee cross-workgroup forward progress, so this mode
8// assumes all NWG workgroups are co-resident (13 workgroups of 128 lanes is
9// far below any real device's residency) and is validated against the
10// multi-dispatch backend, which remains the portable fallback. No thread
11// ever returns early: every guard is non-returning so all threads reach
12// every barrier.
14import { fmt } from './kernel.mjs';
16export function buildPersistent(L, opts = {}) {
17 const wg = opts.wg ?? 128;
18 const chunk = opts.chunk ?? 512;
19 const { nrt, nth, N, nf, K, nsteps, off, A, c } = L;
20 const { ST_CVEL, ST_D, ST_U, ST_UH } = L.layout;
21 const nwg = Math.ceil(N / wg);
23 const constLines = Object.entries(c)
24 .map(([k, v]) => `const ${k}: f32 = ${fmt(v)};`)
25 .join('\n');
27 const code = /* wgsl */ `
28const NRT = ${nrt}u;
29const NTH = ${nth}u;
30const N = ${N}u;
31const NF = ${nf}u;
32const NSTEPS = ${nsteps}u;
33const KS = ${K}u;
34const NRU = ${L.Nr}u;
35const NWG = ${nwg}u;
36const CHUNK = ${chunk}u;
37const LS0 = ${L.LS0}u; const JS0 = ${L.JS0}u;
38const LP0 = ${L.LP0}u; const JP0 = ${L.JP0}u;
39const TIPI = ${L.TIPI0}u;
40const NLR0 = ${L.NLR0}u;
41const HOLDSTEPS = ${L.HOLDSTEPS};
42${constLines}
43${Object.entries(off).map(([k, v]) => `const OP_${k} = ${v}u;`).join('\n')}
44const ST_CVEL = ${ST_CVEL}u;
45const ST_D = ${ST_D}u;
46const ST_U = ${ST_U}u;
47const ST_UH = ${ST_UH}u;
48${L.scratchNames.map((n) => `const A_${n} = ${A[n]}u;`).join('\n')}
50@group(0) @binding(0) var<storage, read> ops: array<f32>;
51@group(0) @binding(1) var<storage, read_write> st: array<f32>;
52@group(0) @binding(2) var<storage, read_write> outb: array<f32>;
53@group(0) @binding(3) var<storage, read_write> sync: array<atomic<u32>, 2>;
55var<workgroup> wgF: f32;
56var<workgroup> wgHA: f32;
57var<workgroup> wgHB: f32;
58var<workgroup> wgHC: f32;
59var<workgroup> shD: array<f32, ${nf}>;
60var<workgroup> shD1: array<f32, ${nf}>;
62fn thm(j: u32) -> u32 { return (j + NTH - 1u) % NTH; }
63fn thp(j: u32) -> u32 { return (j + 1u) % NTH; }
65var<private> P_lid: u32;
67// Classic two-phase global barrier: sync[0] arrive counter, sync[1] epoch.
68fn globalBarrier() {
69 workgroupBarrier();
70 storageBarrier();
71 if (P_lid == 0u) {
72 let gen = atomicLoad(&sync[1]);
73 let n = atomicAdd(&sync[0], 1u);
74 if (n == NWG - 1u) {
75 atomicStore(&sync[0], 0u);
76 atomicAdd(&sync[1], 1u);
77 } else {
78 loop {
79 if (atomicLoad(&sync[1]) != gen) { break; }
80 }
81 }
82 }
83 workgroupBarrier();
84 storageBarrier();
85}
87fn u_dr(b: u32, l: u32, j: u32) -> f32 {
88 if (l < 1u || l > NRT - 2u) { return 0.0; }
89 return (st[b + (l + 1u) + j * NRT] - st[b + (l - 1u) + j * NRT]) * C1R;
90}
92fn p_dr(l: u32, j: u32) -> f32 {
93 if (l < 1u || l > NRT - 2u) { return 0.0; }
94 return (st[A_PHI + (l + 1u) + j * NRT] - st[A_PHI + (l - 1u) + j * NRT]) * C1R;
95}
97@compute @workgroup_size(${wg})
98fn main(@builtin(global_invocation_id) gid: vec3<u32>,
99 @builtin(local_invocation_index) lid: u32,
100 @builtin(workgroup_id) wgid: vec3<u32>) {
101 P_lid = lid;
102 let i = gid.x;
103 let l = i % NRT;
104 let j = i / NRT;
105 let onGrid = i < N;
106 let kbase = u32(st[0u]);
108 for (var s: u32 = 0u; s < CHUNK; s = s + 1u) {
109 let k = kbase + s;
110 let live = k < NSTEPS;
111 let uS = ST_U + (k % 3u) * N;
113 // ---- AB: hand + contact (wg0 lane0), beam (wg0), stencils (all) ----
114 if (wgid.x == 0u) {
115 if (lid == 0u && live) {
116 var shellA = u32(st[1u]); var inC = u32(st[2u]); var phase = u32(st[3u]);
117 var strokeIdx = u32(st[4u]); var holdLeft = i32(st[5u]); var vtar = st[6u];
118 var vH = st[7u]; var uHm = st[8u]; var uH0 = st[9u]; var uHp = st[10u];
119 var peakXi = st[12u]; var peakDisp = st[13u];
120 let dCur = ST_D + ((k + 2u) % 3u) * NF;
121 let dOld = ST_D + ((k + 1u) % 3u) * NF;
122 if (shellA != 0u && k > 0u && (k % 512u) == 0u) {
123 var pd: f32 = 0.0;
124 for (var l2: u32 = 0u; l2 < NRT; l2 = l2 + 1u) {
125 pd = max(pd, abs(st[uS + l2 + JS0 * NRT]));
126 }
127 if (pd > peakDisp) { peakDisp = pd; }
128 }
129 let xi = st[uS + LS0 + JS0 * NRT] - st[dCur + TIPI] - GAP;
130 var fM: f32 = 0.0;
131 if (xi > 0.0) {
132 fM = KM * pow(xi, ALM);
133 shellA = 1u;
134 if (inC == 0u && strokeIdx < KS && st[ST_CVEL + strokeIdx] == 0.0) {
135 st[ST_CVEL + strokeIdx] = (st[dCur + TIPI] - st[dOld + TIPI]) / HT;
136 }
137 inC = 1u;
138 if (xi > peakXi) { peakXi = xi; }
139 } else {
140 inC = 0u;
141 }
142 outb[NSTEPS + k] = fM;
143 wgF = fM;
144 wgHA = -(uHp - 2.0 * uH0 + uHm);
145 wgHB = -(HT * 0.5) * (uHp - uHm);
146 wgHC = -HT * HT * (PS * uH0 + ((1.0 - PS) * 0.5) * (uHp + uHm));
147 if (phase == 0u) {
148 if (strokeIdx < KS && f32(k) >= ops[OP_KGO + strokeIdx]) {
149 phase = 1u; vtar = ops[OP_VSTR + strokeIdx];
150 }
151 } else if (phase == 1u) {
152 vH = max(vH - AH * HT, -vtar);
153 if (inC != 0u) { phase = 3u; holdLeft = HOLDSTEPS; }
154 else if (vH <= -vtar) { phase = 2u; }
155 } else if (phase == 2u) {
156 if (inC != 0u) { phase = 3u; holdLeft = HOLDSTEPS; }
157 else if (uHp <= -MISSD) { phase = 4u; }
158 } else if (phase == 3u) {
159 holdLeft = holdLeft - 1;
160 if (holdLeft <= 0) { phase = 4u; }
161 } else if (phase == 4u) {
162 vH = min(vH + AHR * HT, 0.5 * vtar);
163 if (vH > 0.0 && uHp >= -vH * vH / (2.0 * AHR)) { phase = 5u; }
164 } else {
165 vH = vH - AHR * HT;
166 if (vH <= 0.0 || uHp >= 0.0) { vH = 0.0; phase = 0u; strokeIdx = strokeIdx + 1u; }
167 }
168 uHm = uH0; uH0 = uHp; uHp = uHp + vH * HT;
169 st[1u] = f32(shellA); st[2u] = f32(inC); st[3u] = f32(phase);
170 st[4u] = f32(strokeIdx); st[5u] = f32(holdLeft); st[6u] = vtar;
171 st[7u] = vH; st[8u] = uHm; st[9u] = uH0; st[10u] = uHp;
172 st[12u] = peakXi; st[13u] = peakDisp; st[14u] = fM;
173 }
174 if (lid < NF && live) {
175 shD[lid] = st[ST_D + ((k + 2u) % 3u) * NF + lid];
176 shD1[lid] = st[ST_D + ((k + 1u) % 3u) * NF + lid];
177 }
178 workgroupBarrier();
179 if (lid < NF && live) {
180 let r = lid;
181 var sum: f32 = ops[OP_VM + r] * wgHA + ops[OP_VR + r] * wgHB + ops[OP_VK + r] * wgHC
182 + (HT * HT * wgF) * ops[OP_TIPCOL + r];
183 for (var p: u32 = 0u; p < NF; p = p + 1u) {
184 sum = sum + ops[OP_ASIBT + p + r * NF] * shD[p]
185 + ops[OP_ASICT + p + r * NF] * shD1[p];
186 }
187 st[ST_D + (k % 3u) * NF + r] = sum;
188 }
189 } else {
190 workgroupBarrier(); // match wg0's barrier count
191 }
192 // The shell-active latch for THIS step is decided by wg0 above; other
193 // workgroups may see the pre-latch zero here, which only skips work on a
194 // still-zero field, exactly as in the multi-dispatch backend.
195 var act = onGrid && live && st[1u] != 0.0;
196 if (act) {
197 var ur: f32 = 0.0; var urr: f32 = 0.0;
198 if (l >= 1u && l <= NRT - 2u) {
199 let um = st[uS + (l - 1u) + j * NRT];
200 let u0 = st[uS + l + j * NRT];
201 let up = st[uS + (l + 1u) + j * NRT];
202 ur = (up - um) * C1R;
203 urr = (um - 2.0 * u0 + up) * C2R;
204 }
205 let ujm = st[uS + l + thm(j) * NRT];
206 let uj0 = st[uS + l + j * NRT];
207 let ujp = st[uS + l + thp(j) * NRT];
208 let v1 = ops[OP_RD + l] * ur + ops[OP_R2D + l] * ((ujm - 2.0 * uj0 + ujp) * C2T);
209 let urt = (u_dr(uS, l, thp(j)) - u_dr(uS, l, thm(j))) * C1T;
210 let v2 = ops[OP_RD + l] * urt - ops[OP_R2D + l] * ((ujp - ujm) * C1T);
211 st[A_URR + i] = urr;
212 st[A_V1 + i] = v1;
213 st[A_V2 + i] = v2;
214 var rv: f32 = 0.0;
215 if (l >= 3u && l <= NRU - 2u) {
216 rv = EHRS * (urr + v1) - EH2 * 2.0 * (urr * v1 - v2 * v2);
217 }
218 st[A_RHSA + i] = rv;
219 }
220 globalBarrier();
221 // After the barrier every workgroup sees the latch for this step.
222 act = onGrid && live && st[1u] != 0.0;
224 // ---- C: RHSAF = RHSA * F ----
225 if (act) {
226 var sum: f32 = 0.0;
227 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
228 sum = sum + st[A_RHSA + l + q * NRT] * ops[OP_F + q + j * NTH];
229 }
230 st[A_RHSAF + i] = sum;
231 }
232 globalBarrier();
234 // ---- D: Airy solve per wavenumber ----
235 if (act) {
236 let base = OP_FI + j * NRT * NRT;
237 var sum: f32 = 0.0;
238 for (var p: u32 = 0u; p < NRT; p = p + 1u) {
239 sum = sum + ops[base + l + p * NRT] * st[A_RHSAF + p + j * NRT];
240 }
241 st[A_PH3 + i] = sum;
242 }
243 globalBarrier();
245 // ---- E: PHI = PH3 * Gi ----
246 if (act) {
247 var sum: f32 = 0.0;
248 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
249 sum = sum + st[A_PH3 + l + q * NRT] * ops[OP_GI + q + j * NTH];
250 }
251 st[A_PHI + i] = sum;
252 }
253 globalBarrier();
255 // ---- FG: nonlinear load ----
256 if (act) {
257 var nl: f32 = 0.0;
258 if (l >= NLR0 && l < NRU) {
259 let pm = st[A_PHI + (l - 1u) + j * NRT];
260 let p0 = st[A_PHI + l + j * NRT];
261 let pp = st[A_PHI + (l + 1u) + j * NRT];
262 let pr = (pp - pm) * C1R;
263 let prr = (pm - 2.0 * p0 + pp) * C2R;
264 let pjm = st[A_PHI + l + thm(j) * NRT];
265 let pjp = st[A_PHI + l + thp(j) * NRT];
266 let w1 = ops[OP_RD + l] * pr + ops[OP_R2D + l] * ((pjm - 2.0 * p0 + pjp) * C2T);
267 let prt = (p_dr(l, thp(j)) - p_dr(l, thm(j))) * C1T;
268 let w2 = ops[OP_RD + l] * prt - ops[OP_R2D + l] * ((pjp - pjm) * C1T);
269 nl = prr * st[A_V1 + i] + st[A_URR + i] * w1 - 2.0 * w2 * st[A_V2 + i]
270 - (prr + w1) * RSI;
271 if (l == LS0 && j == JS0) { nl = nl - st[14u] * FSC; }
272 nl = nl * SNL;
273 }
274 st[A_NL + i] = nl;
275 }
276 globalBarrier();
278 // ---- H: NLF = NL * F ----
279 if (act) {
280 var sum: f32 = 0.0;
281 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
282 sum = sum + st[A_NL + l + q * NRT] * ops[OP_F + q + j * NTH];
283 }
284 st[A_NLF + i] = sum;
285 }
286 globalBarrier();
288 // ---- I: implicit update per wavenumber ----
289 if (act) {
290 let h3S = ST_UH + (k % 3u) * N;
291 let h13S = ST_UH + ((k + 2u) % 3u) * N;
292 let h3nS = ST_UH + ((k + 1u) % 3u) * N;
293 let base = j * NRT * NRT;
294 var sum: f32 = 0.0;
295 for (var p: u32 = 0u; p < NRT; p = p + 1u) {
296 let pj = p + j * NRT;
297 sum = sum + ops[OP_AIB + base + l + p * NRT] * st[h3S + pj]
298 + ops[OP_AIC + base + l + p * NRT] * st[h13S + pj]
299 + ops[OP_AI + base + l + p * NRT] * st[A_NLF + pj];
300 }
301 st[h3nS + i] = sum;
302 }
303 globalBarrier();
305 // ---- J: synthesis, pickup acceleration ----
306 if (act) {
307 let h3nS = ST_UH + ((k + 1u) % 3u) * N;
308 let u1S = ST_U + ((k + 2u) % 3u) * N;
309 let unS = ST_U + ((k + 1u) % 3u) * N;
310 var sum: f32 = 0.0;
311 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
312 sum = sum + st[h3nS + l + q * NRT] * ops[OP_GI + q + j * NTH];
313 }
314 st[unS + i] = sum;
315 if (l == LP0 && j == JP0) {
316 outb[k] = (sum - 2.0 * st[uS + i] + st[u1S + i]) * HT2I;
317 }
318 }
319 globalBarrier();
320 }
322 if (wgid.x == 0u && lid == 0u) {
323 st[0u] = f32(min(kbase + CHUNK, NSTEPS));
324 }
325}
326`;
328 return { code, chunk, wg, nwg };
329}