4f822e1turing-surface-cache: reaction-diffusion solutions at a chosen end time, shared through a cloud cacheJeremy Magland 1/**
2 * WGSL Legendre-transform kernels, modeled on leg_m_kernel / ileg_m_kernel
3 * in SHT/cuda_legendre.gen.cu (non-Ishioka fp32 path: SHTNS disables the
4 * Ishioka recurrence for fp32 because it loses too much accuracy).
5 *
6 * Synthesis: F_m(theta_i) = sum_{l=m..lmax} Q_lm * ytilde_l^m(theta_i)
7 * - one thread per latitude, one workgroup row per m (workgroup_id.y).
8 * Analysis: Q_lm = sum_i w_i * G_m(theta_i) * ytilde_l^m(theta_i)
9 * - one workgroup per m; threads own latitudes (strided); per-l pair
10 * workgroup tree reduction (portable stand-in for the CUDA warp
11 * shuffles).
12 *
13 * The associated Legendre functions are generated on the fly by the
14 * standard 3-term recurrence over l (coefficients a,b precomputed on the
15 * host in f64), with the SHTNS fp32 rescaling scheme for sin(theta)^m
16 * underflow (see common.ts).
17 */
18import { RESCALE_WGSL } from './common.ts';
20export interface LegParams {
21 lmax: number;
22 mmax: number;
23 nlat: number;
24 wgSynth: number; // workgroup size for synthesis (threads over latitude)
25 wgAnalys: number; // workgroup size for analysis (power of two)
26 /** Use subgroup reductions in the analysis kernel (needs the `subgroups` feature). */
27 subgroups?: boolean;
28 /** l-pairs accumulated before the span is reduced (subgroup path only). */
29 spanPairs?: number;
30 /**
31 * Fold north/south latitude pairs onto one recurrence (halves Legendre work).
32 * Needs an equator-symmetric grid with even nlat, which the Gauss grid is.
33 */
34 parity?: boolean;
35}
37const BINDINGS = /* wgsl */ `
38@group(0) @binding(0) var<storage, read> ab: array<vec2f>; // (a_l^m, b_l^m) per lm
39@group(0) @binding(1) var<storage, read> amm: array<f32>; // seed per m
40@group(0) @binding(2) var<storage, read> ctstw: array<f32>; // [ct | st | w], each NLAT
41`;
43export function legSynthWGSL(p: LegParams): string {
44 const half = p.parity === true;
45 return /* wgsl */ `
46${RESCALE_WGSL}
47const LMAX: u32 = ${p.lmax}u;
48const NLAT: u32 = ${p.nlat}u;
49const NLAT_2: u32 = ${p.nlat / 2}u;
50${BINDINGS}
51@group(0) @binding(3) var<storage, read> qlm: array<vec2f>;
52@group(0) @binding(4) var<storage, read_write> fm: array<vec2f>; // [(m)*NLAT + ilat]
54@compute @workgroup_size(${p.wgSynth})
55fn leg_synth(@builtin(global_invocation_id) gid: vec3u,
56 @builtin(workgroup_id) wid: vec3u) {
57 let ilat = gid.x;
58 let m = wid.y;
59 if (ilat >= ${half ? 'NLAT_2' : 'NLAT'}) { return; }
61 let ct = ctstw[ilat];
62 let st = ctstw[NLAT + ilat];
63 let base = m * (LMAX + 1u) - (m * (m - 1u)) / 2u; // lm index of (l=m, m)
65 var seed = sinpow_rescaled(st, m);
66 var y0 = seed.y0 * amm[m];
67 var ny = seed.ny;
68 var y1: f32 = 0.0;
69 if (m < LMAX) {
70 y1 = ab[base + 1u].x * ct * y0;
71 }
73${
74 half
75 ? ` // Parity folding: ytilde_l^m(-x) = (-1)^(l-m) ytilde_l^m(x) and the Gauss
76 // grid is symmetric, so one recurrence serves a north/south pair. y0 always
77 // carries even (l-m) and y1 odd, so summing them apart gives
78 // F_m(north) = accE + accO, F_m(south) = accE - accO.
79 var accE = vec2f(0.0);
80 var accO = vec2f(0.0);`
81 : ` var acc = vec2f(0.0);`
82 }
83 var l = m;
84 loop {
85 if (ny == 0) {
86${
87 half
88 ? ` accE += y0 * qlm[base + (l - m)];
89 if (l + 1u <= LMAX) {
90 accO += y1 * qlm[base + (l + 1u - m)];
91 }`
92 : ` acc += y0 * qlm[base + (l - m)];
93 if (l + 1u <= LMAX) {
94 acc += y1 * qlm[base + (l + 1u - m)];
95 }`
96 }
97 } else if (abs(y0) > RESCALE_THR) {
98 ny += 1;
99 y0 *= INV_SCALE;
100 y1 *= INV_SCALE;
101 }
102 if (l + 2u > LMAX) { break; }
103 // Advance (y_l, y_{l+1}) to (y_{l+2}, y_{l+3}).
104 //
105 // Written in exactly the shape leg_analys uses below — both coefficients
106 // fetched unconditionally, the new y0 carried in a temporary rather than
107 // assigned and then read back by the y1 update. The shorter form,
108 //
109 // let c0 = ab[base + (l + 2u - m)];
110 // y0 = c0.x * ct * y1 + c0.y * y0;
111 // if (l + 3u <= LMAX) { ... y1 = c1.x * ct * y0 + c1.y * y1; }
112 //
113 // says the same thing and is what this was, but NVIDIA's Vulkan compiler
114 // (driver 590.48, Blackwell) mis-compiles it: c0 reads as (0, 0) on the
115 // first iteration, so y_{l+2} comes out exactly zero and every later term
116 // follows a different solution of the recurrence, reaching ~1e11 by l = 63.
117 // leg_analys, doing the same arithmetic in this shape, was correct on the
118 // same driver. See scripts/diagnose-leg.ts, which is how that was found.
119 let a0 = ab[base + (l + 2u - m)];
120 var a1 = vec2f(0.0);
121 if (l + 3u <= LMAX) {
122 a1 = ab[base + (l + 3u - m)];
123 }
124 let t0 = a0.x * ct * y1 + a0.y * y0;
125 y1 = a1.x * ct * t0 + a1.y * y1;
126 y0 = t0;
127 l += 2u;
128 }
129${
130 half
131 ? ` fm[m * NLAT + ilat] = accE + accO;
132 fm[m * NLAT + (NLAT - 1u - ilat)] = accE - accO;`
133 : ` fm[m * NLAT + ilat] = acc;`
134 }
135}
136`;
137}
139export function legAnalysWGSL(p: LegParams): string {
140 const half = p.parity === true;
141 // parity folding leaves only the northern half of the grid to walk
142 const K = Math.ceil((half ? p.nlat / 2 : p.nlat) / p.wgAnalys);
143 // With subgroups, the per-l-pair reduction is one subgroupAdd plus a combine
144 // across subgroups: 2 barriers instead of 1 + log2(wgAnalys). This is what
145 // SHTNS's CUDA kernel does with warp shuffles. `red` then holds one partial
146 // per subgroup; WebGPU guarantees subgroup size >= 4, so wgAnalys/4 is a safe
147 // upper bound on how many there can be.
148 const sg = p.subgroups === true;
149 // Reduce once per span of l-pairs rather than once per pair. The l-loop is
150 // serial, so its barriers are the critical path: at lmax=127 the m=0
151 // workgroup paid 2 of them 64 times over. SHTNS amortizes the same way
152 // (LSPAN_A = 16, or 32 for fp32), staging a whole span before reducing.
153 // Partials for the span live in registers and are combined in one batch.
154 const nsubMax = Math.max(1, p.wgAnalys / 4); // WebGPU guarantees subgroup size >= 4
155 // 16 pairs = 32 l-values, which is what SHTNS uses for fp32 (LSPAN_A). Clamped
156 // so `red` stays within 8 KB of workgroup storage, since nsubMax has to assume
157 // the smallest legal subgroup and would otherwise oversize it badly.
158 const pairs = sg
159 ? Math.max(1, Math.min(p.spanPairs ?? 16, Math.floor(8192 / (nsubMax * 16))))
160 : 1;
161 const redLen = sg ? nsubMax * pairs : p.wgAnalys;
162 return /* wgsl */ `${sg ? 'enable subgroups;\n' : ''}
163${RESCALE_WGSL}
164const LMAX: u32 = ${p.lmax}u;
165const NLAT: u32 = ${p.nlat}u;
166const WG: u32 = ${p.wgAnalys}u;
167const K: u32 = ${K}u;
168const NLAT_2: u32 = ${p.nlat / 2}u;
169const PAIRS: u32 = ${pairs}u;
170${BINDINGS}
171@group(0) @binding(3) var<storage, read> fm: array<vec2f>; // [(m)*NLAT + ilat]
172@group(0) @binding(4) var<storage, read_write> qout: array<vec2f>;
174var<workgroup> red: array<vec4f, ${redLen}>;
176@compute @workgroup_size(${p.wgAnalys})
177fn leg_analys(@builtin(local_invocation_id) lid3: vec3u,
178 @builtin(workgroup_id) wid: vec3u${
179 sg
180 ? ',\n @builtin(subgroup_size) sgSize: u32,\n @builtin(subgroup_invocation_id) sgLane: u32'
181 : ''
182 }) {
183 let lid = lid3.x;
184 let m = wid.x;
185 let base = m * (LMAX + 1u) - (m * (m - 1u)) / 2u;
187 // per-thread recurrence state for K latitudes
188 var y0v: array<f32, ${K}>;
189 var y1v: array<f32, ${K}>;
190 var nyv: array<i32, ${K}>;
191 var ctv: array<f32, ${K}>;
192${
193 half
194 ? ` // Transpose of the synthesis folding: splitting the latitude sum into
195 // hemispheres gives Q_lm = sum_north w_i * ytilde * (G_north +/- G_south),
196 // with + for even (l-m) and - for odd -- which the loop already routes
197 // through y0 and y1 respectively.
198 var wpv: array<vec2f, ${K}>;
199 var wmv: array<vec2f, ${K}>;`
200 : ` var wfv: array<vec2f, ${K}>;`
201 }
203 for (var k = 0u; k < K; k++) {
204 let lat = lid + k * WG;
205 var ct: f32 = 0.0;
206 var st: f32 = 0.0;
207${
208 half
209 ? ` var wp = vec2f(0.0);
210 var wm = vec2f(0.0);
211 if (lat < NLAT_2) {
212 ct = ctstw[lat];
213 st = ctstw[NLAT + lat];
214 let w = ctstw[2u * NLAT + lat]; // Gauss weight (incl. 2*pi/nphi)
215 let gN = fm[m * NLAT + lat];
216 let gS = fm[m * NLAT + (NLAT - 1u - lat)];
217 wp = (gN + gS) * w;
218 wm = (gN - gS) * w;
219 }`
220 : ` var wf = vec2f(0.0);
221 if (lat < NLAT) {
222 ct = ctstw[lat];
223 st = ctstw[NLAT + lat];
224 wf = fm[m * NLAT + lat] * ctstw[2u * NLAT + lat]; // Gauss weight (incl. 2*pi/nphi)
225 }`
226 }
227 ctv[k] = ct;
228 let seed = sinpow_rescaled(st, m);
229 y0v[k] = seed.y0 * amm[m];
230 nyv[k] = seed.ny;
231 y1v[k] = 0.0;
232 if (m < LMAX) {
233 y1v[k] = ab[base + 1u].x * ct * y0v[k];
234 }
235${half ? ' wpv[k] = wp;\n wmv[k] = wm;' : ' wfv[k] = wf;'}
236 }
238 var l = m;
239${
240 sg
241 ? ` // Accumulate up to PAIRS l-pairs into registers, then reduce the whole span
242 // at once: 2 barriers per span instead of 2 per pair.
243 loop {
244 let lstart = l;
245 var npairs = 0u;
246 var last = false;
247 let sub = lid / sgSize;
248 for (var jj = 0u; jj < PAIRS; jj++) {
249 var c0 = vec2f(0.0);
250 var c1 = vec2f(0.0);
251 for (var k = 0u; k < K; k++) {
252 if (nyv[k] == 0) {
253${
254 half
255 ? ` c0 += wpv[k] * y0v[k]; // even (l-m): hemispheres add
256 c1 += wmv[k] * y1v[k]; // odd (l-m): hemispheres subtract`
257 : ` c0 += wfv[k] * y0v[k];
258 c1 += wfv[k] * y1v[k];`
259 }
260 } else if (abs(y0v[k]) > RESCALE_THR) {
261 nyv[k] += 1;
262 y0v[k] *= INV_SCALE;
263 y1v[k] *= INV_SCALE;
264 }
265 }
266 // subgroupAdd needs no barrier, so the per-subgroup partial can go
267 // straight to shared memory; only the cross-subgroup combine below has
268 // to wait, and it waits once for the whole span.
269 let part = subgroupAdd(vec4f(c0, c1));
270 if (sgLane == 0u) { red[sub * PAIRS + jj] = part; }
271 npairs = jj + 1u;
272 if (l + 2u > LMAX) { last = true; break; }
273 let a0 = ab[base + (l + 2u - m)];
274 var a1 = vec2f(0.0);
275 if (l + 3u <= LMAX) {
276 a1 = ab[base + (l + 3u - m)];
277 }
278 for (var k = 0u; k < K; k++) {
279 let t0 = a0.x * ctv[k] * y1v[k] + a0.y * y0v[k];
280 y0v[k] = t0;
281 y1v[k] = a1.x * ctv[k] * t0 + a1.y * y1v[k];
282 }
283 l += 2u;
284 }
286 workgroupBarrier();
287 if (lid == 0u) {
288 let nsub = (WG + sgSize - 1u) / sgSize;
289 for (var jj = 0u; jj < npairs; jj++) {
290 var tot = vec4f(0.0);
291 for (var i = 0u; i < nsub; i++) { tot += red[i * PAIRS + jj]; }
292 let ll = lstart + 2u * jj;
293 qout[base + (ll - m)] = tot.xy;
294 if (ll + 1u <= LMAX) {
295 qout[base + (ll + 1u - m)] = tot.zw;
296 }
297 }
298 }
299 workgroupBarrier(); // red is reused by the next span
301 if (last) { break; }
302 }`
303 : ` loop {
304 var c0 = vec2f(0.0);
305 var c1 = vec2f(0.0);
306 for (var k = 0u; k < K; k++) {
307 if (nyv[k] == 0) {
308${
309 half
310 ? ` c0 += wpv[k] * y0v[k]; // even (l-m): hemispheres add
311 c1 += wmv[k] * y1v[k]; // odd (l-m): hemispheres subtract`
312 : ` c0 += wfv[k] * y0v[k];
313 c1 += wfv[k] * y1v[k];`
314 }
315 } else if (abs(y0v[k]) > RESCALE_THR) {
316 nyv[k] += 1;
317 y0v[k] *= INV_SCALE;
318 y1v[k] *= INV_SCALE;
319 }
320 }
321 // workgroup tree reduction of (c0, c1)
322 red[lid] = vec4f(c0, c1);
323 workgroupBarrier();
324 var s = WG / 2u;
325 while (s > 0u) {
326 if (lid < s) { red[lid] += red[lid + s]; }
327 workgroupBarrier();
328 s = s >> 1u;
329 }
330 if (lid == 0u) {
331 qout[base + (l - m)] = red[0].xy;
332 if (l + 1u <= LMAX) {
333 qout[base + (l + 1u - m)] = red[0].zw;
334 }
335 }
336 if (l + 2u > LMAX) { break; }
337 let a0 = ab[base + (l + 2u - m)];
338 var a1 = vec2f(0.0);
339 if (l + 3u <= LMAX) {
340 a1 = ab[base + (l + 3u - m)];
341 }
342 for (var k = 0u; k < K; k++) {
343 let t0 = a0.x * ctv[k] * y1v[k] + a0.y * y0v[k];
344 y0v[k] = t0;
345 y1v[k] = a1.x * ctv[k] * t0 + a1.y * y1v[k];
346 }
347 l += 2u;
348 }`
349 }
350}
351`;
352}
354/**
355 * Batched transforms: K independent fields through ONE walk of the Legendre
356 * recurrence. The recurrence state (y0, y1, rescaling) depends only on
357 * (m, theta), never on the field, so a batch shares it and pays only the
358 * extra data fetches and accumulators per lane — the same amortization
359 * SHTNS's GPU backend gets from batching fields. Per-lane arithmetic is
360 * textually identical to the scalar kernels' (same operations, same order),
361 * so a batched transform reproduces the scalar transform's results.
362 *
363 * K is a codegen parameter. The bind group needs 3 tables + K inputs +
364 * K outputs storage buffers, so K = 2 (7 bindings) fits WebGPU's default
365 * limit of 8 on every stack including SwiftShader, and K = 4 (11) needs the
366 * raised limit requestShtDevice asks for where the adapter offers it.
367 * The recurrence bodies are kept in exactly the shape the scalar kernels
368 * use — see the driver-workaround comment in legSynthWGSL before
369 * "simplifying" either copy.
370 */
371export function legSynthBatchWGSL(p: LegParams, K: number, laneElems: number): string {
372 const half = p.parity === true;
373 const lanes = Array.from({ length: K }, (_, k) => k);
374 // K caller-owned inputs, ONE plan-owned fm arena: lane k writes at a fixed
375 // 256-byte-aligned offset (laneElems vec2f), which is what keeps the bind
376 // group at 3 + K + 1 storage buffers -- within WebGPU's default limit of 8
377 // at K = 4. The Fourier stage binds the arena per lane with a buffer
378 // offset, so it needs no changes.
379 const bind =
380 lanes
381 .map((k) => `@group(0) @binding(${3 + k}) var<storage, read> qlm${k}: array<vec2f>;`)
382 .join('\n') +
383 `\n@group(0) @binding(${3 + K}) var<storage, read_write> fm: array<vec2f>;`;
384 const decl = lanes
385 .map((k) =>
386 half
387 ? ` var accE${k} = vec2f(0.0);\n var accO${k} = vec2f(0.0);`
388 : ` var acc${k} = vec2f(0.0);`,
389 )
390 .join('\n');
391 const accEven = lanes
392 .map((k) => (half ? ` accE${k} += y0 * qlm${k}[i0];` : ` acc${k} += y0 * qlm${k}[i0];`))
393 .join('\n');
394 const accOdd = lanes
395 .map((k) => (half ? ` accO${k} += y1 * qlm${k}[i1];` : ` acc${k} += y1 * qlm${k}[i1];`))
396 .join('\n');
397 const store = lanes
398 .map((k) =>
399 half
400 ? ` fm[${k}u * LANE + m * NLAT + ilat] = accE${k} + accO${k};\n` +
401 ` fm[${k}u * LANE + m * NLAT + (NLAT - 1u - ilat)] = accE${k} - accO${k};`
402 : ` fm[${k}u * LANE + m * NLAT + ilat] = acc${k};`,
403 )
404 .join('\n');
405 return /* wgsl */ `
406${RESCALE_WGSL}
407const LMAX: u32 = ${p.lmax}u;
408const NLAT: u32 = ${p.nlat}u;
409const NLAT_2: u32 = ${p.nlat / 2}u;
410const LANE: u32 = ${laneElems}u;
411${BINDINGS}
412${bind}
414@compute @workgroup_size(${p.wgSynth})
415fn leg_synth_batch(@builtin(global_invocation_id) gid: vec3u,
416 @builtin(workgroup_id) wid: vec3u) {
417 let ilat = gid.x;
418 let m = wid.y;
419 if (ilat >= ${half ? 'NLAT_2' : 'NLAT'}) { return; }
421 let ct = ctstw[ilat];
422 let st = ctstw[NLAT + ilat];
423 let base = m * (LMAX + 1u) - (m * (m - 1u)) / 2u;
425 var seed = sinpow_rescaled(st, m);
426 var y0 = seed.y0 * amm[m];
427 var ny = seed.ny;
428 var y1: f32 = 0.0;
429 if (m < LMAX) {
430 y1 = ab[base + 1u].x * ct * y0;
431 }
433${decl}
434 var l = m;
435 loop {
436 if (ny == 0) {
437 let i0 = base + (l - m);
438${accEven}
439 if (l + 1u <= LMAX) {
440 let i1 = base + (l + 1u - m);
441${accOdd}
442 }
443 } else if (abs(y0) > RESCALE_THR) {
444 ny += 1;
445 y0 *= INV_SCALE;
446 y1 *= INV_SCALE;
447 }
448 if (l + 2u > LMAX) { break; }
449 // Same two-coefficient, temporary-carried shape as leg_synth (see the
450 // driver-workaround comment there).
451 let a0 = ab[base + (l + 2u - m)];
452 var a1 = vec2f(0.0);
453 if (l + 3u <= LMAX) {
454 a1 = ab[base + (l + 3u - m)];
455 }
456 let t0 = a0.x * ct * y1 + a0.y * y0;
457 y1 = a1.x * ct * t0 + a1.y * y1;
458 y0 = t0;
459 l += 2u;
460 }
461${store}
462}
463`;
464}
466/** Batched analysis: K spatial-Fourier fields reduced against one Legendre
467 * recurrence walk. Structure follows legAnalysWGSL; see legSynthBatchWGSL
468 * for the batching rationale and the binding budget. */
469export function legAnalysBatchWGSL(p: LegParams, K: number, laneElems: number): string {
470 const half = p.parity === true;
471 const lanes = Array.from({ length: K }, (_, k) => k);
472 const Kl = Math.ceil((half ? p.nlat / 2 : p.nlat) / p.wgAnalys);
473 const sg = p.subgroups === true;
474 const nsubMax = Math.max(1, p.wgAnalys / 4);
475 // Same 8 KB workgroup-storage budget as the scalar kernel, now split
476 // across K lanes, so spans shorten as K grows: barriers per unit of work
477 // stay level.
478 const pairs = sg
479 ? Math.max(1, Math.min(p.spanPairs ?? 16, Math.floor(8192 / (nsubMax * 16 * K))))
480 : 1;
481 const redLen = (sg ? nsubMax * pairs : p.wgAnalys) * K;
482 // ONE fm arena in (lane offsets baked, as in legSynthBatchWGSL), K
483 // caller-owned outputs: 3 + 1 + K storage buffers.
484 const bind =
485 `@group(0) @binding(3) var<storage, read> fm: array<vec2f>;\n` +
486 lanes
487 .map((k) => `@group(0) @binding(${4 + k}) var<storage, read_write> qout${k}: array<vec2f>;`)
488 .join('\n');
489 const laneState = lanes
490 .map((k) =>
491 half
492 ? ` var wpv${k}: array<vec2f, ${Kl}>;\n var wmv${k}: array<vec2f, ${Kl}>;`
493 : ` var wfv${k}: array<vec2f, ${Kl}>;`,
494 )
495 .join('\n');
496 const laneLoad = lanes
497 .map((k) =>
498 half
499 ? ` let gN${k} = fm[${k}u * LANE + m * NLAT + lat];
500 let gS${k} = fm[${k}u * LANE + m * NLAT + (NLAT - 1u - lat)];
501 wp${k} = (gN${k} + gS${k}) * w;
502 wm${k} = (gN${k} - gS${k}) * w;`
503 : ` wf${k} = fm[${k}u * LANE + m * NLAT + lat] * w;`,
504 )
505 .join('\n');
506 const laneLoadDecl = lanes
507 .map((k) =>
508 half ? ` var wp${k} = vec2f(0.0);\n var wm${k} = vec2f(0.0);` : ` var wf${k} = vec2f(0.0);`,
509 )
510 .join('\n');
511 const laneLoadStore = lanes
512 .map((k) => (half ? ` wpv${k}[k] = wp${k};\n wmv${k}[k] = wm${k};` : ` wfv${k}[k] = wf${k};`))
513 .join('\n');
514 const cDecl = lanes.map((k) => ` var c0_${k} = vec2f(0.0);\n var c1_${k} = vec2f(0.0);`).join('\n');
515 const cAcc = lanes
516 .map((k) =>
517 half
518 ? ` c0_${k} += wpv${k}[k] * y0v[k];
519 c1_${k} += wmv${k}[k] * y1v[k];`
520 : ` c0_${k} += wfv${k}[k] * y0v[k];
521 c1_${k} += wfv${k}[k] * y1v[k];`,
522 )
523 .join('\n');
524 return /* wgsl */ `${sg ? 'enable subgroups;\n' : ''}
525${RESCALE_WGSL}
526const LMAX: u32 = ${p.lmax}u;
527const NLAT: u32 = ${p.nlat}u;
528const WG: u32 = ${p.wgAnalys}u;
529const K: u32 = ${Kl}u;
530const NLAT_2: u32 = ${p.nlat / 2}u;
531const PAIRS: u32 = ${pairs}u;
532const NB: u32 = ${K}u;
533const LANE: u32 = ${laneElems}u;
534${BINDINGS}
535${bind}
537var<workgroup> red: array<vec4f, ${redLen}>;
539@compute @workgroup_size(${p.wgAnalys})
540fn leg_analys_batch(@builtin(local_invocation_id) lid3: vec3u,
541 @builtin(workgroup_id) wid: vec3u${
542 sg
543 ? ',\n @builtin(subgroup_size) sgSize: u32,\n @builtin(subgroup_invocation_id) sgLane: u32'
544 : ''
545 }) {
546 let lid = lid3.x;
547 let m = wid.x;
548 let base = m * (LMAX + 1u) - (m * (m - 1u)) / 2u;
550 var y0v: array<f32, ${Kl}>;
551 var y1v: array<f32, ${Kl}>;
552 var nyv: array<i32, ${Kl}>;
553 var ctv: array<f32, ${Kl}>;
554${laneState}
556 for (var k = 0u; k < K; k++) {
557 let lat = lid + k * WG;
558 var ct: f32 = 0.0;
559 var st: f32 = 0.0;
560 var w: f32 = 0.0;
561${laneLoadDecl}
562 if (lat < ${half ? 'NLAT_2' : 'NLAT'}) {
563 ct = ctstw[lat];
564 st = ctstw[NLAT + lat];
565 w = ctstw[2u * NLAT + lat];
566${laneLoad}
567 }
568 ctv[k] = ct;
569 let seed = sinpow_rescaled(st, m);
570 y0v[k] = seed.y0 * amm[m];
571 nyv[k] = seed.ny;
572 y1v[k] = 0.0;
573 if (m < LMAX) {
574 y1v[k] = ab[base + 1u].x * ct * y0v[k];
575 }
576${laneLoadStore}
577 }
579 var l = m;
580${
581 sg
582 ? ` loop {
583 let lstart = l;
584 var npairs = 0u;
585 var last = false;
586 let sub = lid / sgSize;
587 for (var jj = 0u; jj < PAIRS; jj++) {
588${cDecl}
589 for (var k = 0u; k < K; k++) {
590 if (nyv[k] == 0) {
591${cAcc}
592 } else if (abs(y0v[k]) > RESCALE_THR) {
593 nyv[k] += 1;
594 y0v[k] *= INV_SCALE;
595 y1v[k] *= INV_SCALE;
596 }
597 }
598${lanes
599 .map(
600 (k) => ` let part${k} = subgroupAdd(vec4f(c0_${k}, c1_${k}));
601 if (sgLane == 0u) { red[(sub * PAIRS + jj) * NB + ${k}u] = part${k}; }`,
602 )
603 .join('\n')}
604 npairs = jj + 1u;
605 if (l + 2u > LMAX) { last = true; break; }
606 let a0 = ab[base + (l + 2u - m)];
607 var a1 = vec2f(0.0);
608 if (l + 3u <= LMAX) {
609 a1 = ab[base + (l + 3u - m)];
610 }
611 for (var k = 0u; k < K; k++) {
612 let t0 = a0.x * ctv[k] * y1v[k] + a0.y * y0v[k];
613 y0v[k] = t0;
614 y1v[k] = a1.x * ctv[k] * t0 + a1.y * y1v[k];
615 }
616 l += 2u;
617 }
619 workgroupBarrier();
620 if (lid == 0u) {
621 let nsub = (WG + sgSize - 1u) / sgSize;
622 for (var jj = 0u; jj < npairs; jj++) {
623 let ll = lstart + 2u * jj;
624${lanes
625 .map(
626 (k) => ` var tot${k} = vec4f(0.0);
627 for (var i = 0u; i < nsub; i++) { tot${k} += red[(i * PAIRS + jj) * NB + ${k}u]; }
628 qout${k}[base + (ll - m)] = tot${k}.xy;
629 if (ll + 1u <= LMAX) {
630 qout${k}[base + (ll + 1u - m)] = tot${k}.zw;
631 }`,
632 )
633 .join('\n')}
634 }
635 }
636 workgroupBarrier(); // red is reused by the next span
638 if (last) { break; }
639 }`
640 : ` loop {
641${cDecl}
642 for (var k = 0u; k < K; k++) {
643 if (nyv[k] == 0) {
644${cAcc.replace(/^ {10}/gm, ' ')}
645 } else if (abs(y0v[k]) > RESCALE_THR) {
646 nyv[k] += 1;
647 y0v[k] *= INV_SCALE;
648 y1v[k] *= INV_SCALE;
649 }
650 }
651 // workgroup tree reduction, lane-strided
652${lanes.map((k) => ` red[lid + ${k}u * WG] = vec4f(c0_${k}, c1_${k});`).join('\n')}
653 workgroupBarrier();
654 var s = WG / 2u;
655 while (s > 0u) {
656 if (lid < s) {
657${lanes.map((k) => ` red[lid + ${k}u * WG] += red[lid + s + ${k}u * WG];`).join('\n')}
658 }
659 workgroupBarrier();
660 s = s >> 1u;
661 }
662 if (lid == 0u) {
663${lanes
664 .map(
665 (k) => ` qout${k}[base + (l - m)] = red[${k}u * WG].xy;
666 if (l + 1u <= LMAX) {
667 qout${k}[base + (l + 1u - m)] = red[${k}u * WG].zw;
668 }`,
669 )
670 .join('\n')}
671 }
672 if (l + 2u > LMAX) { break; }
673 let a0 = ab[base + (l + 2u - m)];
674 var a1 = vec2f(0.0);
675 if (l + 3u <= LMAX) {
676 a1 = ab[base + (l + 3u - m)];
677 }
678 for (var k = 0u; k < K; k++) {
679 let t0 = a0.x * ctv[k] * y1v[k] + a0.y * y0v[k];
680 y0v[k] = t0;
681 y1v[k] = a1.x * ctv[k] * t0 + a1.y * y1v[k];
682 }
683 l += 2u;
684 }`
685 }
686}
687`;
688}