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}
30const BINDINGS = /* wgsl */ `
31@group(0) @binding(0) var<storage, read> ab: array<vec2f>; // (a_l^m, b_l^m) per lm
32@group(0) @binding(1) var<storage, read> amm: array<f32>; // seed per m
33@group(0) @binding(2) var<storage, read> ctstw: array<f32>; // [ct | st | w], each NLAT
34`;
36export function legSynthWGSL(p: LegParams): string {
37 return /* wgsl */ `
38${RESCALE_WGSL}
39const LMAX: u32 = ${p.lmax}u;
40const NLAT: u32 = ${p.nlat}u;
41${BINDINGS}
42@group(0) @binding(3) var<storage, read> qlm: array<vec2f>;
43@group(0) @binding(4) var<storage, read_write> fm: array<vec2f>; // [(m)*NLAT + ilat]
45@compute @workgroup_size(${p.wgSynth})
46fn leg_synth(@builtin(global_invocation_id) gid: vec3u,
47 @builtin(workgroup_id) wid: vec3u) {
48 let ilat = gid.x;
49 let m = wid.y;
50 if (ilat >= NLAT) { return; }
52 let ct = ctstw[ilat];
53 let st = ctstw[NLAT + ilat];
54 let base = m * (LMAX + 1u) - (m * (m - 1u)) / 2u; // lm index of (l=m, m)
56 var seed = sinpow_rescaled(st, m);
57 var y0 = seed.y0 * amm[m];
58 var ny = seed.ny;
59 var y1: f32 = 0.0;
60 if (m < LMAX) {
61 y1 = ab[base + 1u].x * ct * y0;
62 }
64 var acc = vec2f(0.0);
65 var l = m;
66 loop {
67 if (ny == 0) {
68 acc += y0 * qlm[base + (l - m)];
69 if (l + 1u <= LMAX) {
70 acc += y1 * qlm[base + (l + 1u - m)];
71 }
72 } else if (abs(y0) > RESCALE_THR) {
73 ny += 1;
74 y0 *= INV_SCALE;
75 y1 *= INV_SCALE;
76 }
77 if (l + 2u > LMAX) { break; }
78 // Advance (y_l, y_{l+1}) to (y_{l+2}, y_{l+3}).
79 //
80 // Written in exactly the shape leg_analys uses below — both coefficients
81 // fetched unconditionally, the new y0 carried in a temporary rather than
82 // assigned and then read back by the y1 update. The shorter form,
83 //
84 // let c0 = ab[base + (l + 2u - m)];
85 // y0 = c0.x * ct * y1 + c0.y * y0;
86 // if (l + 3u <= LMAX) { ... y1 = c1.x * ct * y0 + c1.y * y1; }
87 //
88 // says the same thing and is what this was, but NVIDIA's Vulkan compiler
89 // (driver 590.48, Blackwell) mis-compiles it: c0 reads as (0, 0) on the
90 // first iteration, so y_{l+2} comes out exactly zero and every later term
91 // follows a different solution of the recurrence, reaching ~1e11 by l = 63.
92 // leg_analys, doing the same arithmetic in this shape, was correct on the
93 // same driver. See scripts/diagnose-leg.ts, which is how that was found.
94 let a0 = ab[base + (l + 2u - m)];
95 var a1 = vec2f(0.0);
96 if (l + 3u <= LMAX) {
97 a1 = ab[base + (l + 3u - m)];
98 }
99 let t0 = a0.x * ct * y1 + a0.y * y0;
100 y1 = a1.x * ct * t0 + a1.y * y1;
101 y0 = t0;
102 l += 2u;
103 }
104 fm[m * NLAT + ilat] = acc;
105}
106`;
107}
109export function legAnalysWGSL(p: LegParams): string {
110 const K = Math.ceil(p.nlat / p.wgAnalys); // latitudes per thread
111 // With subgroups, the per-l-pair reduction is one subgroupAdd plus a combine
112 // across subgroups: 2 barriers instead of 1 + log2(wgAnalys). This is what
113 // SHTNS's CUDA kernel does with warp shuffles. `red` then holds one partial
114 // per subgroup; WebGPU guarantees subgroup size >= 4, so wgAnalys/4 is a safe
115 // upper bound on how many there can be.
116 const sg = p.subgroups === true;
117 const redLen = sg ? Math.max(1, p.wgAnalys / 4) : p.wgAnalys;
118 return /* wgsl */ `${sg ? 'enable subgroups;\n' : ''}
119${RESCALE_WGSL}
120const LMAX: u32 = ${p.lmax}u;
121const NLAT: u32 = ${p.nlat}u;
122const WG: u32 = ${p.wgAnalys}u;
123const K: u32 = ${K}u;
124${BINDINGS}
125@group(0) @binding(3) var<storage, read> fm: array<vec2f>; // [(m)*NLAT + ilat]
126@group(0) @binding(4) var<storage, read_write> qout: array<vec2f>;
128var<workgroup> red: array<vec4f, ${redLen}>;
130@compute @workgroup_size(${p.wgAnalys})
131fn leg_analys(@builtin(local_invocation_id) lid3: vec3u,
132 @builtin(workgroup_id) wid: vec3u${
133 sg
134 ? ',\n @builtin(subgroup_size) sgSize: u32,\n @builtin(subgroup_invocation_id) sgLane: u32'
135 : ''
136 }) {
137 let lid = lid3.x;
138 let m = wid.x;
139 let base = m * (LMAX + 1u) - (m * (m - 1u)) / 2u;
141 // per-thread recurrence state for K latitudes
142 var y0v: array<f32, ${K}>;
143 var y1v: array<f32, ${K}>;
144 var nyv: array<i32, ${K}>;
145 var ctv: array<f32, ${K}>;
146 var wfv: array<vec2f, ${K}>;
148 for (var k = 0u; k < K; k++) {
149 let lat = lid + k * WG;
150 var ct: f32 = 0.0;
151 var st: f32 = 0.0;
152 var wf = vec2f(0.0);
153 if (lat < NLAT) {
154 ct = ctstw[lat];
155 st = ctstw[NLAT + lat];
156 wf = fm[m * NLAT + lat] * ctstw[2u * NLAT + lat]; // Gauss weight (incl. 2*pi/nphi)
157 }
158 ctv[k] = ct;
159 let seed = sinpow_rescaled(st, m);
160 y0v[k] = seed.y0 * amm[m];
161 nyv[k] = seed.ny;
162 y1v[k] = 0.0;
163 if (m < LMAX) {
164 y1v[k] = ab[base + 1u].x * ct * y0v[k];
165 }
166 wfv[k] = wf;
167 }
169 var l = m;
170 loop {
171 var c0 = vec2f(0.0);
172 var c1 = vec2f(0.0);
173 for (var k = 0u; k < K; k++) {
174 if (nyv[k] == 0) {
175 c0 += wfv[k] * y0v[k];
176 c1 += wfv[k] * y1v[k];
177 } else if (abs(y0v[k]) > RESCALE_THR) {
178 nyv[k] += 1;
179 y0v[k] *= INV_SCALE;
180 y1v[k] *= INV_SCALE;
181 }
182 }
183${
184 sg
185 ? ` // Reduce (c0, c1) across the workgroup: one subgroupAdd, then combine the
186 // per-subgroup partials. Two barriers per l-pair rather than 1 + log2(WG),
187 // and no half-idle tree. Both barriers sit in uniform control flow.
188 let part = subgroupAdd(vec4f(c0, c1));
189 if (sgLane == 0u) { red[lid / sgSize] = part; }
190 workgroupBarrier();
191 if (lid == 0u) {
192 var tot = vec4f(0.0);
193 let nsub = (WG + sgSize - 1u) / sgSize;
194 for (var i = 0u; i < nsub; i++) { tot += red[i]; }
195 qout[base + (l - m)] = tot.xy;
196 if (l + 1u <= LMAX) {
197 qout[base + (l + 1u - m)] = tot.zw;
198 }
199 }
200 workgroupBarrier(); // red is reused next iteration`
201 : ` // workgroup tree reduction of (c0, c1)
202 red[lid] = vec4f(c0, c1);
203 workgroupBarrier();
204 var s = WG / 2u;
205 while (s > 0u) {
206 if (lid < s) { red[lid] += red[lid + s]; }
207 workgroupBarrier();
208 s = s >> 1u;
209 }
210 if (lid == 0u) {
211 qout[base + (l - m)] = red[0].xy;
212 if (l + 1u <= LMAX) {
213 qout[base + (l + 1u - m)] = red[0].zw;
214 }
215 }`
216 }
217 if (l + 2u > LMAX) { break; }
218 let a0 = ab[base + (l + 2u - m)];
219 var a1 = vec2f(0.0);
220 if (l + 3u <= LMAX) {
221 a1 = ab[base + (l + 3u - m)];
222 }
223 for (var k = 0u; k < K; k++) {
224 let t0 = a0.x * ctv[k] * y1v[k] + a0.y * y0v[k];
225 y0v[k] = t0;
226 y1v[k] = a1.x * ctv[k] * t0 + a1.y * y1v[k];
227 }
228 l += 2u;
229 }
230}
231`;
232}