2dedc35turing-sphere: reaction-diffusion on the sphere, spectral solver on WebGPUJeremy 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)
53ff7c5leg_analys: reduce with subgroupAdd, and retune wgAnalys for itdanfortunato 26 /** Use subgroup reductions in the analysis kernel (needs the `subgroups` feature). */
27 subgroups?: boolean;
b81424bleg_analys: reduce once per span of l-pairs, not once per pairdanfortunato 28 /** l-pairs accumulated before the span is reduced (subgroup path only). */
29 spanPairs?: number;
32const BINDINGS = /* wgsl */ `
33@group(0) @binding(0) var<storage, read> ab: array<vec2f>; // (a_l^m, b_l^m) per lm
34@group(0) @binding(1) var<storage, read> amm: array<f32>; // seed per m
35@group(0) @binding(2) var<storage, read> ctstw: array<f32>; // [ct | st | w], each NLAT
36`;
38export function legSynthWGSL(p: LegParams): string {
39 return /* wgsl */ `
40${RESCALE_WGSL}
41const LMAX: u32 = ${p.lmax}u;
42const NLAT: u32 = ${p.nlat}u;
43${BINDINGS}
44@group(0) @binding(3) var<storage, read> qlm: array<vec2f>;
45@group(0) @binding(4) var<storage, read_write> fm: array<vec2f>; // [(m)*NLAT + ilat]
47@compute @workgroup_size(${p.wgSynth})
48fn leg_synth(@builtin(global_invocation_id) gid: vec3u,
49 @builtin(workgroup_id) wid: vec3u) {
50 let ilat = gid.x;
51 let m = wid.y;
52 if (ilat >= NLAT) { return; }
54 let ct = ctstw[ilat];
55 let st = ctstw[NLAT + ilat];
56 let base = m * (LMAX + 1u) - (m * (m - 1u)) / 2u; // lm index of (l=m, m)
58 var seed = sinpow_rescaled(st, m);
59 var y0 = seed.y0 * amm[m];
60 var ny = seed.ny;
61 var y1: f32 = 0.0;
62 if (m < LMAX) {
63 y1 = ab[base + 1u].x * ct * y0;
64 }
66 var acc = vec2f(0.0);
67 var l = m;
68 loop {
69 if (ny == 0) {
70 acc += y0 * qlm[base + (l - m)];
71 if (l + 1u <= LMAX) {
72 acc += y1 * qlm[base + (l + 1u - m)];
73 }
74 } else if (abs(y0) > RESCALE_THR) {
75 ny += 1;
76 y0 *= INV_SCALE;
77 y1 *= INV_SCALE;
78 }
79 if (l + 2u > LMAX) { break; }
f290310leg_synth: advance the recurrence the way leg_analys doesJeremy Magland 80 // Advance (y_l, y_{l+1}) to (y_{l+2}, y_{l+3}).
81 //
82 // Written in exactly the shape leg_analys uses below — both coefficients
83 // fetched unconditionally, the new y0 carried in a temporary rather than
84 // assigned and then read back by the y1 update. The shorter form,
85 //
86 // let c0 = ab[base + (l + 2u - m)];
87 // y0 = c0.x * ct * y1 + c0.y * y0;
88 // if (l + 3u <= LMAX) { ... y1 = c1.x * ct * y0 + c1.y * y1; }
89 //
90 // says the same thing and is what this was, but NVIDIA's Vulkan compiler
91 // (driver 590.48, Blackwell) mis-compiles it: c0 reads as (0, 0) on the
92 // first iteration, so y_{l+2} comes out exactly zero and every later term
93 // follows a different solution of the recurrence, reaching ~1e11 by l = 63.
94 // leg_analys, doing the same arithmetic in this shape, was correct on the
95 // same driver. See scripts/diagnose-leg.ts, which is how that was found.
96 let a0 = ab[base + (l + 2u - m)];
97 var a1 = vec2f(0.0);
2dedc35turing-sphere: reaction-diffusion on the sphere, spectral solver on WebGPUJeremy Magland 98 if (l + 3u <= LMAX) {
f290310leg_synth: advance the recurrence the way leg_analys doesJeremy Magland 99 a1 = ab[base + (l + 3u - m)];
2dedc35turing-sphere: reaction-diffusion on the sphere, spectral solver on WebGPUJeremy Magland 100 }
f290310leg_synth: advance the recurrence the way leg_analys doesJeremy Magland 101 let t0 = a0.x * ct * y1 + a0.y * y0;
102 y1 = a1.x * ct * t0 + a1.y * y1;
103 y0 = t0;
2dedc35turing-sphere: reaction-diffusion on the sphere, spectral solver on WebGPUJeremy Magland 104 l += 2u;
105 }
106 fm[m * NLAT + ilat] = acc;
107}
108`;
109}
111export function legAnalysWGSL(p: LegParams): string {
112 const K = Math.ceil(p.nlat / p.wgAnalys); // latitudes per thread
53ff7c5leg_analys: reduce with subgroupAdd, and retune wgAnalys for itdanfortunato 113 // With subgroups, the per-l-pair reduction is one subgroupAdd plus a combine
114 // across subgroups: 2 barriers instead of 1 + log2(wgAnalys). This is what
115 // SHTNS's CUDA kernel does with warp shuffles. `red` then holds one partial
116 // per subgroup; WebGPU guarantees subgroup size >= 4, so wgAnalys/4 is a safe
117 // upper bound on how many there can be.
118 const sg = p.subgroups === true;
b81424bleg_analys: reduce once per span of l-pairs, not once per pairdanfortunato 119 // Reduce once per span of l-pairs rather than once per pair. The l-loop is
120 // serial, so its barriers are the critical path: at lmax=127 the m=0
121 // workgroup paid 2 of them 64 times over. SHTNS amortizes the same way
122 // (LSPAN_A = 16, or 32 for fp32), staging a whole span before reducing.
123 // Partials for the span live in registers and are combined in one batch.
124 const nsubMax = Math.max(1, p.wgAnalys / 4); // WebGPU guarantees subgroup size >= 4
125 // 16 pairs = 32 l-values, which is what SHTNS uses for fp32 (LSPAN_A). Clamped
126 // so `red` stays within 8 KB of workgroup storage, since nsubMax has to assume
127 // the smallest legal subgroup and would otherwise oversize it badly.
128 const pairs = sg
129 ? Math.max(1, Math.min(p.spanPairs ?? 16, Math.floor(8192 / (nsubMax * 16))))
130 : 1;
131 const redLen = sg ? nsubMax * pairs : p.wgAnalys;
53ff7c5leg_analys: reduce with subgroupAdd, and retune wgAnalys for itdanfortunato 132 return /* wgsl */ `${sg ? 'enable subgroups;\n' : ''}
2dedc35turing-sphere: reaction-diffusion on the sphere, spectral solver on WebGPUJeremy Magland 133${RESCALE_WGSL}
134const LMAX: u32 = ${p.lmax}u;
135const NLAT: u32 = ${p.nlat}u;
136const WG: u32 = ${p.wgAnalys}u;
137const K: u32 = ${K}u;
b81424bleg_analys: reduce once per span of l-pairs, not once per pairdanfortunato 138const PAIRS: u32 = ${pairs}u;
2dedc35turing-sphere: reaction-diffusion on the sphere, spectral solver on WebGPUJeremy Magland 139${BINDINGS}
140@group(0) @binding(3) var<storage, read> fm: array<vec2f>; // [(m)*NLAT + ilat]
141@group(0) @binding(4) var<storage, read_write> qout: array<vec2f>;
53ff7c5leg_analys: reduce with subgroupAdd, and retune wgAnalys for itdanfortunato 143var<workgroup> red: array<vec4f, ${redLen}>;
145@compute @workgroup_size(${p.wgAnalys})
146fn leg_analys(@builtin(local_invocation_id) lid3: vec3u,
53ff7c5leg_analys: reduce with subgroupAdd, and retune wgAnalys for itdanfortunato 147 @builtin(workgroup_id) wid: vec3u${
148 sg
149 ? ',\n @builtin(subgroup_size) sgSize: u32,\n @builtin(subgroup_invocation_id) sgLane: u32'
150 : ''
151 }) {
2dedc35turing-sphere: reaction-diffusion on the sphere, spectral solver on WebGPUJeremy Magland 152 let lid = lid3.x;
153 let m = wid.x;
154 let base = m * (LMAX + 1u) - (m * (m - 1u)) / 2u;
156 // per-thread recurrence state for K latitudes
157 var y0v: array<f32, ${K}>;
158 var y1v: array<f32, ${K}>;
159 var nyv: array<i32, ${K}>;
160 var ctv: array<f32, ${K}>;
161 var wfv: array<vec2f, ${K}>;
163 for (var k = 0u; k < K; k++) {
164 let lat = lid + k * WG;
165 var ct: f32 = 0.0;
166 var st: f32 = 0.0;
167 var wf = vec2f(0.0);
168 if (lat < NLAT) {
169 ct = ctstw[lat];
170 st = ctstw[NLAT + lat];
171 wf = fm[m * NLAT + lat] * ctstw[2u * NLAT + lat]; // Gauss weight (incl. 2*pi/nphi)
172 }
173 ctv[k] = ct;
174 let seed = sinpow_rescaled(st, m);
175 y0v[k] = seed.y0 * amm[m];
176 nyv[k] = seed.ny;
177 y1v[k] = 0.0;
178 if (m < LMAX) {
179 y1v[k] = ab[base + 1u].x * ct * y0v[k];
180 }
181 wfv[k] = wf;
182 }
184 var l = m;
186 sg
187 ? ` // Accumulate up to PAIRS l-pairs into registers, then reduce the whole span
188 // at once: 2 barriers per span instead of 2 per pair.
2dedc35turing-sphere: reaction-diffusion on the sphere, spectral solver on WebGPUJeremy Magland 189 loop {
b81424bleg_analys: reduce once per span of l-pairs, not once per pairdanfortunato 190 let lstart = l;
191 var npairs = 0u;
192 var last = false;
193 let sub = lid / sgSize;
194 for (var jj = 0u; jj < PAIRS; jj++) {
195 var c0 = vec2f(0.0);
196 var c1 = vec2f(0.0);
197 for (var k = 0u; k < K; k++) {
198 if (nyv[k] == 0) {
199 c0 += wfv[k] * y0v[k];
200 c1 += wfv[k] * y1v[k];
201 } else if (abs(y0v[k]) > RESCALE_THR) {
202 nyv[k] += 1;
203 y0v[k] *= INV_SCALE;
204 y1v[k] *= INV_SCALE;
205 }
206 }
207 // subgroupAdd needs no barrier, so the per-subgroup partial can go
208 // straight to shared memory; only the cross-subgroup combine below has
209 // to wait, and it waits once for the whole span.
210 let part = subgroupAdd(vec4f(c0, c1));
211 if (sgLane == 0u) { red[sub * PAIRS + jj] = part; }
212 npairs = jj + 1u;
213 if (l + 2u > LMAX) { last = true; break; }
214 let a0 = ab[base + (l + 2u - m)];
215 var a1 = vec2f(0.0);
216 if (l + 3u <= LMAX) {
217 a1 = ab[base + (l + 3u - m)];
218 }
219 for (var k = 0u; k < K; k++) {
220 let t0 = a0.x * ctv[k] * y1v[k] + a0.y * y0v[k];
221 y0v[k] = t0;
222 y1v[k] = a1.x * ctv[k] * t0 + a1.y * y1v[k];
223 }
224 l += 2u;
225 }
227 workgroupBarrier();
228 if (lid == 0u) {
229 let nsub = (WG + sgSize - 1u) / sgSize;
230 for (var jj = 0u; jj < npairs; jj++) {
231 var tot = vec4f(0.0);
232 for (var i = 0u; i < nsub; i++) { tot += red[i * PAIRS + jj]; }
233 let ll = lstart + 2u * jj;
234 qout[base + (ll - m)] = tot.xy;
235 if (ll + 1u <= LMAX) {
236 qout[base + (ll + 1u - m)] = tot.zw;
237 }
238 }
239 }
240 workgroupBarrier(); // red is reused by the next span
242 if (last) { break; }
243 }`
244 : ` loop {
2dedc35turing-sphere: reaction-diffusion on the sphere, spectral solver on WebGPUJeremy Magland 245 var c0 = vec2f(0.0);
246 var c1 = vec2f(0.0);
247 for (var k = 0u; k < K; k++) {
248 if (nyv[k] == 0) {
249 c0 += wfv[k] * y0v[k];
250 c1 += wfv[k] * y1v[k];
251 } else if (abs(y0v[k]) > RESCALE_THR) {
252 nyv[k] += 1;
253 y0v[k] *= INV_SCALE;
254 y1v[k] *= INV_SCALE;
255 }
256 }
b81424bleg_analys: reduce once per span of l-pairs, not once per pairdanfortunato 257 // workgroup tree reduction of (c0, c1)
2dedc35turing-sphere: reaction-diffusion on the sphere, spectral solver on WebGPUJeremy Magland 258 red[lid] = vec4f(c0, c1);
259 workgroupBarrier();
260 var s = WG / 2u;
261 while (s > 0u) {
262 if (lid < s) { red[lid] += red[lid + s]; }
263 workgroupBarrier();
264 s = s >> 1u;
265 }
266 if (lid == 0u) {
267 qout[base + (l - m)] = red[0].xy;
268 if (l + 1u <= LMAX) {
269 qout[base + (l + 1u - m)] = red[0].zw;
270 }
2dedc35turing-sphere: reaction-diffusion on the sphere, spectral solver on WebGPUJeremy Magland 272 if (l + 2u > LMAX) { 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;
2dedc35turing-sphere: reaction-diffusion on the sphere, spectral solver on WebGPUJeremy Magland 285 }
286}
287`;
288}