1// Multi-dispatch backend: each stage of a time step is its own compute
2// dispatch at full GPU width, and one step is 8 dispatches:
3// AB hand state machine + contact + beam (workgroup 0), and the shell
4// derivative stencils / Airy right-hand side (all workgroups)
5// C RHSAF = RHSA * F (azimuthal analysis)
6// D PH3 = Fi_j RHSAF_j (Airy solve per wavenumber)
7// E PHI = PH3 * Gi (azimuthal synthesis)
8// FG NL from PHI (stress-function stencils recomputed inline), forcing
9// H NLF = NL * F
10// I UH3N = AiB UH3 + AiC UH13 + Ai NLF (implicit update per column)
11// J UNEW = UH3N * Gi, pickup acceleration
12//
13// The step index k is kbase (state cell 0, written by the host once per
14// chunk) plus s from a dynamic-offset uniform, so every recorded dispatch is
15// identical up to that offset and slot rotation is a pure function of k.
16// U/Uh slots rotate every step whether or not the shell is active: before
17// first contact all slots are zero, so the rotation is harmless, and after
18// first contact every step writes. The shell-active flag latches inside AB;
19// stages of the same step read it a dispatch later, so they always see the
20// latch, and the benign race inside AB itself can only cause zero-work on a
21// zero field.
23import { fmt } from './kernel.mjs';
25export function buildStages(L, opts = {}) {
26 const wg = opts.wg ?? Number(process.env.WG ?? 64);
27 const { nrt, nth, N, nf, K, nsteps, off, A, c } = L;
28 const { SC, ST_CVEL, ST_D, ST_U, ST_UH } = L.layout;
30 const constLines = Object.entries(c)
31 .map(([k, v]) => `const ${k}: f32 = ${fmt(v)};`)
32 .join('\n');
34 const header = `
35const NRT = ${nrt}u;
36const NTH = ${nth}u;
37const N = ${N}u;
38const NF = ${nf}u;
39const NSTEPS = ${nsteps}u;
40const KS = ${K}u;
41const NRU = ${L.Nr}u;
42const LS0 = ${L.LS0}u; const JS0 = ${L.JS0}u;
43const LP0 = ${L.LP0}u; const JP0 = ${L.JP0}u;
44const TIPI = ${L.TIPI0}u;
45const NLR0 = ${L.NLR0}u;
46const HOLDSTEPS = ${L.HOLDSTEPS};
47${constLines}
48${Object.entries(off).map(([k, v]) => `const OP_${k} = ${v}u;`).join('\n')}
49const ST_CVEL = ${ST_CVEL}u;
50const ST_D = ${ST_D}u;
51const ST_U = ${ST_U}u;
52const ST_UH = ${ST_UH}u;
53${L.scratchNames.map((n) => `const A_${n} = ${A[n]}u;`).join('\n')}
55@group(0) @binding(0) var<storage, read> ops: array<f32>;
56@group(0) @binding(1) var<storage, read_write> st: array<f32>;
57@group(0) @binding(2) var<storage, read_write> outb: array<f32>;
58@group(0) @binding(3) var<uniform> ctl: vec4<u32>;
60fn thm(j: u32) -> u32 { return (j + NTH - 1u) % NTH; }
61fn thp(j: u32) -> u32 { return (j + 1u) % NTH; }
62`;
65 const stages = {};
67 stages.AB = `${header}
68var<workgroup> wgF: f32;
69var<workgroup> wgHA: f32;
70var<workgroup> wgHB: f32;
71var<workgroup> wgHC: f32;
72var<workgroup> shD: array<f32, ${nf}>;
73var<workgroup> shD1: array<f32, ${nf}>;
75fn u_dr(b: u32, l: u32, j: u32) -> f32 {
76 if (l < 1u || l > NRT - 2u) { return 0.0; }
77 return (st[b + (l + 1u) + j * NRT] - st[b + (l - 1u) + j * NRT]) * C1R;
78}
80@compute @workgroup_size(${wg})
81fn main(@builtin(global_invocation_id) gid: vec3<u32>,
82 @builtin(local_invocation_index) lid: u32,
83 @builtin(workgroup_id) wgid: vec3<u32>) {
84 // No early-out on k here: the host records exactly nsteps steps, and the
85 // workgroupBarrier below requires uniform control flow up to this point.
86 let k = u32(st[0u]) + ctl.x;
87 let uS = ST_U + (k % 3u) * N;
88 if (wgid.x == 0u) {
89 if (lid == 0u) {
90 var shellA = u32(st[1u]); var inC = u32(st[2u]); var phase = u32(st[3u]);
91 var strokeIdx = u32(st[4u]); var holdLeft = i32(st[5u]); var vtar = st[6u];
92 var vH = st[7u]; var uHm = st[8u]; var uH0 = st[9u]; var uHp = st[10u];
93 var peakXi = st[12u]; var peakDisp = st[13u];
94 let dCur = ST_D + ((k + 2u) % 3u) * NF;
95 let dOld = ST_D + ((k + 1u) % 3u) * NF;
96 if (shellA != 0u && k > 0u && (k % 512u) == 0u) {
97 var pd: f32 = 0.0;
98 for (var l: u32 = 0u; l < NRT; l = l + 1u) {
99 pd = max(pd, abs(st[uS + l + JS0 * NRT]));
100 }
101 if (pd > peakDisp) { peakDisp = pd; }
102 }
103 let xi = st[uS + LS0 + JS0 * NRT] - st[dCur + TIPI] - GAP;
104 var fM: f32 = 0.0;
105 if (xi > 0.0) {
106 fM = KM * pow(xi, ALM);
107 shellA = 1u;
108 if (inC == 0u && strokeIdx < KS && st[ST_CVEL + strokeIdx] == 0.0) {
109 st[ST_CVEL + strokeIdx] = (st[dCur + TIPI] - st[dOld + TIPI]) / HT;
110 }
111 inC = 1u;
112 if (xi > peakXi) { peakXi = xi; }
113 } else {
114 inC = 0u;
115 }
116 outb[NSTEPS + k] = fM;
117 wgF = fM;
118 wgHA = -(uHp - 2.0 * uH0 + uHm);
119 wgHB = -(HT * 0.5) * (uHp - uHm);
120 wgHC = -HT * HT * (PS * uH0 + ((1.0 - PS) * 0.5) * (uHp + uHm));
121 if (phase == 0u) {
122 if (strokeIdx < KS && f32(k) >= ops[OP_KGO + strokeIdx]) {
123 phase = 1u; vtar = ops[OP_VSTR + strokeIdx];
124 }
125 } else if (phase == 1u) {
126 vH = max(vH - AH * HT, -vtar);
127 if (inC != 0u) { phase = 3u; holdLeft = HOLDSTEPS; }
128 else if (vH <= -vtar) { phase = 2u; }
129 } else if (phase == 2u) {
130 if (inC != 0u) { phase = 3u; holdLeft = HOLDSTEPS; }
131 else if (uHp <= -MISSD) { phase = 4u; }
132 } else if (phase == 3u) {
133 holdLeft = holdLeft - 1;
134 if (holdLeft <= 0) { phase = 4u; }
135 } else if (phase == 4u) {
136 vH = min(vH + AHR * HT, 0.5 * vtar);
137 if (vH > 0.0 && uHp >= -vH * vH / (2.0 * AHR)) { phase = 5u; }
138 } else {
139 vH = vH - AHR * HT;
140 if (vH <= 0.0 || uHp >= 0.0) { vH = 0.0; phase = 0u; strokeIdx = strokeIdx + 1u; }
141 }
142 uHm = uH0; uH0 = uHp; uHp = uHp + vH * HT;
143 st[1u] = f32(shellA); st[2u] = f32(inC); st[3u] = f32(phase);
144 st[4u] = f32(strokeIdx); st[5u] = f32(holdLeft); st[6u] = vtar;
145 st[7u] = vH; st[8u] = uHm; st[9u] = uH0; st[10u] = uHp;
146 st[12u] = peakXi; st[13u] = peakDisp; st[14u] = fM;
147 }
148 if (lid < NF) {
149 shD[lid] = st[ST_D + ((k + 2u) % 3u) * NF + lid];
150 shD1[lid] = st[ST_D + ((k + 1u) % 3u) * NF + lid];
151 }
152 workgroupBarrier();
153 if (lid < NF) {
154 let r = lid;
155 var sum: f32 = ops[OP_VM + r] * wgHA + ops[OP_VR + r] * wgHB + ops[OP_VK + r] * wgHC
156 + (HT * HT * wgF) * ops[OP_TIPCOL + r];
157 for (var p: u32 = 0u; p < NF; p = p + 1u) {
158 sum = sum + ops[OP_ASIBT + p + r * NF] * shD[p]
159 + ops[OP_ASICT + p + r * NF] * shD1[p];
160 }
161 st[ST_D + (k % 3u) * NF + r] = sum;
162 }
163 }
164 let i = gid.x;
165 if (i >= N) { return; }
166 if (st[1u] == 0.0) { return; }
167 let l = i % NRT; let j = i / NRT;
168 var ur: f32 = 0.0; var urr: f32 = 0.0;
169 if (l >= 1u && l <= NRT - 2u) {
170 let um = st[uS + (l - 1u) + j * NRT];
171 let u0 = st[uS + l + j * NRT];
172 let up = st[uS + (l + 1u) + j * NRT];
173 ur = (up - um) * C1R;
174 urr = (um - 2.0 * u0 + up) * C2R;
175 }
176 let ujm = st[uS + l + thm(j) * NRT];
177 let uj0 = st[uS + l + j * NRT];
178 let ujp = st[uS + l + thp(j) * NRT];
179 let v1 = ops[OP_RD + l] * ur + ops[OP_R2D + l] * ((ujm - 2.0 * uj0 + ujp) * C2T);
180 let urt = (u_dr(uS, l, thp(j)) - u_dr(uS, l, thm(j))) * C1T;
181 let v2 = ops[OP_RD + l] * urt - ops[OP_R2D + l] * ((ujp - ujm) * C1T);
182 st[A_URR + i] = urr;
183 st[A_V1 + i] = v1;
184 st[A_V2 + i] = v2;
185 var rv: f32 = 0.0;
186 if (l >= 3u && l <= NRU - 2u) {
187 rv = EHRS * (urr + v1) - EH2 * 2.0 * (urr * v1 - v2 * v2);
188 }
189 st[A_RHSA + i] = rv;
190}
191`;
193 const gate = `
194 let k = u32(st[0u]) + ctl.x;
195 if (k >= NSTEPS) { return; }
196 let i = gid.x;
197 if (i >= N) { return; }
198 if (st[1u] == 0.0) { return; }
199 let l = i % NRT; let j = i / NRT;`;
201 const transform = (src, dst, mat) => `${header}
202@compute @workgroup_size(${wg})
203fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
204${gate}
205 var sum: f32 = 0.0;
206 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
207 sum = sum + st[${src} + l + q * NRT] * ops[${mat} + q + j * NTH];
208 }
209 st[${dst} + i] = sum;
210}
211`;
213 stages.C = transform('A_RHSA', 'A_RHSAF', 'OP_F');
215 stages.D = `${header}
216@compute @workgroup_size(${wg})
217fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
218${gate}
219 let base = OP_FI + j * NRT * NRT;
220 var sum: f32 = 0.0;
221 for (var p: u32 = 0u; p < NRT; p = p + 1u) {
222 sum = sum + ops[base + l + p * NRT] * st[A_RHSAF + p + j * NRT];
223 }
224 st[A_PH3 + i] = sum;
225}
226`;
228 stages.E = transform('A_PH3', 'A_PHI', 'OP_GI');
230 stages.FG = `${header}
231fn p_dr(l: u32, j: u32) -> f32 {
232 if (l < 1u || l > NRT - 2u) { return 0.0; }
233 return (st[A_PHI + (l + 1u) + j * NRT] - st[A_PHI + (l - 1u) + j * NRT]) * C1R;
234}
235@compute @workgroup_size(${wg})
236fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
237${gate}
238 var nl: f32 = 0.0;
239 if (l >= NLR0 && l < NRU) {
240 let pm = st[A_PHI + (l - 1u) + j * NRT];
241 let p0 = st[A_PHI + l + j * NRT];
242 let pp = st[A_PHI + (l + 1u) + j * NRT];
243 let pr = (pp - pm) * C1R;
244 let prr = (pm - 2.0 * p0 + pp) * C2R;
245 let pjm = st[A_PHI + l + thm(j) * NRT];
246 let pjp = st[A_PHI + l + thp(j) * NRT];
247 let w1 = ops[OP_RD + l] * pr + ops[OP_R2D + l] * ((pjm - 2.0 * p0 + pjp) * C2T);
248 let prt = (p_dr(l, thp(j)) - p_dr(l, thm(j))) * C1T;
249 let w2 = ops[OP_RD + l] * prt - ops[OP_R2D + l] * ((pjp - pjm) * C1T);
250 nl = prr * st[A_V1 + i] + st[A_URR + i] * w1 - 2.0 * w2 * st[A_V2 + i]
251 - (prr + w1) * RSI;
252 if (l == LS0 && j == JS0) { nl = nl - st[14u] * FSC; }
253 nl = nl * SNL;
254 }
255 st[A_NL + i] = nl;
256}
257`;
259 stages.H = transform('A_NL', 'A_NLF', 'OP_F');
261 stages.I = `${header}
262@compute @workgroup_size(${wg})
263fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
264${gate}
265 let h3S = ST_UH + (k % 3u) * N;
266 let h13S = ST_UH + ((k + 2u) % 3u) * N;
267 let h3nS = ST_UH + ((k + 1u) % 3u) * N;
268 let base = j * NRT * NRT;
269 var sum: f32 = 0.0;
270 for (var p: u32 = 0u; p < NRT; p = p + 1u) {
271 let pj = p + j * NRT;
272 sum = sum + ops[OP_AIB + base + l + p * NRT] * st[h3S + pj]
273 + ops[OP_AIC + base + l + p * NRT] * st[h13S + pj]
274 + ops[OP_AI + base + l + p * NRT] * st[A_NLF + pj];
275 }
276 st[h3nS + i] = sum;
277}
278`;
280 stages.J = `${header}
281@compute @workgroup_size(${wg})
282fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
283${gate}
284 let h3nS = ST_UH + ((k + 1u) % 3u) * N;
285 let uS = ST_U + (k % 3u) * N;
286 let u1S = ST_U + ((k + 2u) % 3u) * N;
287 let unS = ST_U + ((k + 1u) % 3u) * N;
288 var sum: f32 = 0.0;
289 for (var q: u32 = 0u; q < NTH; q = q + 1u) {
290 sum = sum + st[h3nS + l + q * NRT] * ops[OP_GI + q + j * NTH];
291 }
292 st[unS + i] = sum;
293 if (l == LP0 && j == JP0) {
294 outb[k] = (sum - 2.0 * st[uS + i] + st[u1S + i]) * HT2I;
295 }
296}
297`;
299 const nwgMain = Math.ceil(N / wg);
300 return {
301 order: ['AB', 'C', 'D', 'E', 'FG', 'H', 'I', 'J'],
302 sources: stages,
303 wg,
304 nwg: nwgMain,
305 };
306}