Fold north/south latitude pairs onto one Legendre recurrence - concept-collection/turing-sphere

concept-collection / turing-sphere

Fold north/south latitude pairs onto one Legendre recurrence

The Gauss grid is symmetric about the equator and ytilde_l^m(-x) = (-1)^(l-m) ytilde_l^m(x), so one recurrence serves a north/south pair. Both kernels already separate the two parities -- y0 carries even (l-m), y1 odd -- so synthesis needs only to accumulate them apart, F_m(north) = accE + accO and F_m(south) = accE - accO, and walk half the latitudes. Analysis is the transpose: split the latitude sum into hemispheres and each northern latitude contributes w_i * ytilde * (G_north +/- G_south), + for even (l-m) and - for odd. SHTNS does the same. The saving is in K = ceil(nlat/wgAnalys), the latitudes each analysis thread carries, which parity halves. So the gain appears only once K > 1: grid K off K on round trip off -> on 128x256 1 1 34.0 -> 34.0 1.00x 256x512 2 1 60.6 -> 58.6 1.03x 512x1024 4 2 118.7 -> 110.0 1.08x npm run bench --preset allencahn: lmax=127 10733 -> 11158 steps/s, lmax=255 6285 -> 6925. Round-trip accuracy unchanged at 3.87e-6 relative L2; test:node and diagnose-leg pass. Worth recording that this was tried before the two reduction commits and measured a no-op at 128x256 and 12% *worse* at 512x1024, which is why it was reverted then. With a barrier on every l-pair the kernels were reduction-bound, so halving the recurrence work bought nothing and the extra wpv/wmv state made it negative. Once the reduction is amortized over a span the arithmetic saving surfaces and the regression is gone. The ordering matters: this is only worth applying on top of the span-batched reduction.

danfortunato <dan.fortunato@gmail.com> committed 20 days ago commit eb4a9e73778d parent b81424b Browse files

2 changed files+78−13

src/sht/sht.tsmodified+7−1View file

@@ -115,6 +115,8 @@ export class ShtPlan {

115115 readonly cfg: ShtConfig;

116116 readonly nlm: number;

117117 readonly fourierMode: 'fft' | 'dft';

118+ /** Latitudes leg_synth walks: nlat/2 when parity folding. */

119+ readonly legLat: number = 0;

118120 /** Colatitudes theta_i (f64, increasing: north to south). */

119121 readonly theta: Float64Array;

120122 readonly cosTheta: Float64Array;

@@ -220,6 +222,8 @@ export class ShtPlan {

220222

221223 // --- shaders / pipelines ---

222224 const subgroups = tuning('SHT_SUBGROUPS') !== false && dev.features.has('subgroups');

225+ // parity folding needs an equator-symmetric grid; Gauss nodes are, if nlat is even

226+ const parity = tuning('SHT_PARITY') !== false && nlat % 2 === 0;

223227 const wgAnalys =

224228 (tuning('SHT_WG_ANALYS') as number | undefined) ??

225229 defaultWgAnalys(nlat, dev.limits.maxComputeInvocationsPerWorkgroup, subgroups);

@@ -231,7 +235,9 @@ export class ShtPlan {

231235 wgAnalys,

232236 subgroups,

233237 spanPairs: tuning('SHT_SPAN_PAIRS') as number | undefined,

238+ parity,

234239 };

240+ (this as { legLat: number }).legLat = parity ? nlat / 2 : nlat;

235241 const fourP = { mmax, nlat, nphi };

236242 const [pLegS, pLegA, pFourS, pFourA] = await Promise.all([

237243 makePipeline(dev, legSynthWGSL(legP), 'leg_synth'),

@@ -310,7 +316,7 @@ export class ShtPlan {

310316 const { mmax, nlat, nphi } = this.cfg;

311317 pass.setPipeline(this.pipeLegSynth);

312318 pass.setBindGroup(0, b.bgLeg);

313- pass.dispatchWorkgroups(Math.ceil(nlat / WG_SYNTH), mmax + 1);

319+ pass.dispatchWorkgroups(Math.ceil(this.legLat / WG_SYNTH), mmax + 1);

314320 pass.setPipeline(this.pipeFourSynth);

315321 pass.setBindGroup(0, b.bgFour);

316322 if (this.fourierMode === 'fft') {

src/sht/wgsl/leg.tsmodified+71−12View file

@@ -27,6 +27,11 @@ export interface LegParams {

2727 subgroups?: boolean;

2828 /** l-pairs accumulated before the span is reduced (subgroup path only). */

2929 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;

3035 }

3136

3237 const BINDINGS = /* wgsl */ `

@@ -36,10 +41,12 @@ const BINDINGS = /* wgsl */ `

3641 `;

3742

3843 export function legSynthWGSL(p: LegParams): string {

44+ const half = p.parity === true;

3945 return /* wgsl */ `

4046 ${RESCALE_WGSL}

4147 const LMAX: u32 = ${p.lmax}u;

4248 const NLAT: u32 = ${p.nlat}u;

49+const NLAT_2: u32 = ${p.nlat / 2}u;

4350 ${BINDINGS}

4451 @group(0) @binding(3) var<storage, read> qlm: array<vec2f>;

4552 @group(0) @binding(4) var<storage, read_write> fm: array<vec2f>; // [(m)*NLAT + ilat]

@@ -49,7 +56,7 @@ fn leg_synth(@builtin(global_invocation_id) gid: vec3u,

4956 @builtin(workgroup_id) wid: vec3u) {

5057 let ilat = gid.x;

5158 let m = wid.y;

52- if (ilat >= NLAT) { return; }

59+ if (ilat >= ${half ? 'NLAT_2' : 'NLAT'}) { return; }

5360

5461 let ct = ctstw[ilat];

5562 let st = ctstw[NLAT + ilat];

@@ -63,14 +70,30 @@ fn leg_synth(@builtin(global_invocation_id) gid: vec3u,

6370 y1 = ab[base + 1u].x * ct * y0;

6471 }

6572

66- var acc = vec2f(0.0);

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+ }

6783 var l = m;

6884 loop {

6985 if (ny == 0) {

70- acc += y0 * qlm[base + (l - m)];

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)];

7193 if (l + 1u <= LMAX) {

7294 acc += y1 * qlm[base + (l + 1u - m)];

73- }

95+ }`

96+ }

7497 } else if (abs(y0) > RESCALE_THR) {

7598 ny += 1;

7699 y0 *= INV_SCALE;

@@ -103,13 +126,20 @@ fn leg_synth(@builtin(global_invocation_id) gid: vec3u,

103126 y0 = t0;

104127 l += 2u;

105128 }

106- fm[m * NLAT + ilat] = acc;

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+ }

107135 }

108136 `;

109137 }

110138

111139 export function legAnalysWGSL(p: LegParams): string {

112- const K = Math.ceil(p.nlat / p.wgAnalys); // latitudes per thread

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);

113143 // With subgroups, the per-l-pair reduction is one subgroupAdd plus a combine

114144 // across subgroups: 2 barriers instead of 1 + log2(wgAnalys). This is what

115145 // SHTNS's CUDA kernel does with warp shuffles. `red` then holds one partial

@@ -135,6 +165,7 @@ const LMAX: u32 = ${p.lmax}u;

135165 const NLAT: u32 = ${p.nlat}u;

136166 const WG: u32 = ${p.wgAnalys}u;

137167 const K: u32 = ${K}u;

168+const NLAT_2: u32 = ${p.nlat / 2}u;

138169 const PAIRS: u32 = ${pairs}u;

139170 ${BINDINGS}

140171 @group(0) @binding(3) var<storage, read> fm: array<vec2f>; // [(m)*NLAT + ilat]

@@ -158,18 +189,41 @@ fn leg_analys(@builtin(local_invocation_id) lid3: vec3u,

158189 var y1v: array<f32, ${K}>;

159190 var nyv: array<i32, ${K}>;

160191 var ctv: array<f32, ${K}>;

161- var wfv: array<vec2f, ${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+ }

162202

163203 for (var k = 0u; k < K; k++) {

164204 let lat = lid + k * WG;

165205 var ct: f32 = 0.0;

166206 var st: f32 = 0.0;

167- var wf = vec2f(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);

168221 if (lat < NLAT) {

169222 ct = ctstw[lat];

170223 st = ctstw[NLAT + lat];

171224 wf = fm[m * NLAT + lat] * ctstw[2u * NLAT + lat]; // Gauss weight (incl. 2*pi/nphi)

172- }

225+ }`

226+ }

173227 ctv[k] = ct;

174228 let seed = sinpow_rescaled(st, m);

175229 y0v[k] = seed.y0 * amm[m];

@@ -178,7 +232,7 @@ fn leg_analys(@builtin(local_invocation_id) lid3: vec3u,

178232 if (m < LMAX) {

179233 y1v[k] = ab[base + 1u].x * ct * y0v[k];

180234 }

181- wfv[k] = wf;

235+${half ? ' wpv[k] = wp;\n wmv[k] = wm;' : ' wfv[k] = wf;'}

182236 }

183237

184238 var l = m;

@@ -196,8 +250,13 @@ ${

196250 var c1 = vec2f(0.0);

197251 for (var k = 0u; k < K; k++) {

198252 if (nyv[k] == 0) {

199- c0 += wfv[k] * y0v[k];

200- c1 += wfv[k] * y1v[k];

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+ }

201260 } else if (abs(y0v[k]) > RESCALE_THR) {

202261 nyv[k] += 1;

203262 y0v[k] *= INV_SCALE;