leg_analys: reduce once per span of l-pairs, not once per pair - concept-collection/turing-sphere

concept-collection / turing-sphere

leg_analys: reduce once per span of l-pairs, not once per pair

The l-loop is serial, so the two workgroup barriers around the cross-subgroup combine sat on the critical path once per l-pair: at lmax=127 the m=0 workgroup paid them 64 times over. SHTNS amortizes the same cost with LSPAN_A (16, or 32 for fp32), staging a whole span of l before reducing it -- one barrier per span rather than per coefficient. subgroupAdd needs no barrier, so a per-subgroup partial can go straight to shared memory as each pair is finished; only the cross-subgroup combine has to wait, and it now waits once for the whole span. Barriers per l-pair go from 2 to 2/PAIRS. The recurrence and thread layout are untouched. PAIRS defaults to 16, i.e. 32 l-values, matching SHTNS's fp32 LSPAN_A, and is clamped so `red` stays inside 8 KB: nsubMax has to assume the smallest subgroup size WebGPU guarantees (4), so it oversizes the array otherwise. SHT_SPAN_PAIRS overrides it. Measured on an RTX PRO 6000 Blackwell, scripts/_split.ts (300 round trips per submit, best of 3), us: analysis round trip grid before after before after 128x256 22.1 17.9 38.3 34.0 1.13x 256x512 40.5 34.3 66.6 60.5 1.10x 512x1024 80.3 67.5 131.0 119.3 1.10x Analysis and synthesis are now near parity (17.9 vs 16.0 us at lmax=63), where analysis had been 3.3x synthesis before any of this work -- the reduction overhead is essentially gone. What remains in the round trip is ~15 us of fixed dispatch cost across four dispatches. npm run bench --preset allencahn --lmax 127: 9960 -> 10900 steps/s. Cumulative over the two commits, 6090 -> 10900, 1.79x. Round-trip accuracy unchanged at 3.87e-6 relative L2 (lmax=127); test:node and diagnose-leg pass. Note for anyone extending this: staging the span in a register array (var acc: array<vec4f, PAIRS>) is slower than not batching at all -- the dynamic index spills it to scratch, and PAIRS=1 then measured 0.106 against 0.079 ms for the unbatched code. Writing partials straight to shared memory avoids the array entirely.

danfortunato <dan.fortunato@gmail.com> committed 20 days ago commit b81424b713e5 parent 53ff7c5 Browse files

2 changed files+79−22

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

@@ -230,6 +230,7 @@ export class ShtPlan {

230230 wgSynth: WG_SYNTH,

231231 wgAnalys,

232232 subgroups,

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

233234 };

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

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

src/sht/wgsl/leg.tsmodified+78−22View file

@@ -25,6 +25,8 @@ export interface LegParams {

2525 wgAnalys: number; // workgroup size for analysis (power of two)

2626 /** Use subgroup reductions in the analysis kernel (needs the `subgroups` feature). */

2727 subgroups?: boolean;

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

29+ spanPairs?: number;

2830 }

2931

3032 const BINDINGS = /* wgsl */ `

@@ -114,13 +116,26 @@ export function legAnalysWGSL(p: LegParams): string {

114116 // per subgroup; WebGPU guarantees subgroup size >= 4, so wgAnalys/4 is a safe

115117 // upper bound on how many there can be.

116118 const sg = p.subgroups === true;

117- const redLen = sg ? Math.max(1, p.wgAnalys / 4) : p.wgAnalys;

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;

118132 return /* wgsl */ `${sg ? 'enable subgroups;\n' : ''}

119133 ${RESCALE_WGSL}

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

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

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

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

138+const PAIRS: u32 = ${pairs}u;

124139 ${BINDINGS}

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

126141 @group(0) @binding(4) var<storage, read_write> qout: array<vec2f>;

@@ -167,7 +182,66 @@ fn leg_analys(@builtin(local_invocation_id) lid3: vec3u,

167182 }

168183

169184 var l = m;

185+${

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.

170189 loop {

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

226+

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

241+

242+ if (last) { break; }

243+ }`

244+ : ` loop {

171245 var c0 = vec2f(0.0);

172246 var c1 = vec2f(0.0);

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

@@ -180,25 +254,7 @@ fn leg_analys(@builtin(local_invocation_id) lid3: vec3u,

180254 y1v[k] *= INV_SCALE;

181255 }

182256 }

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)

257+ // workgroup tree reduction of (c0, c1)

202258 red[lid] = vec4f(c0, c1);

203259 workgroupBarrier();

204260 var s = WG / 2u;

@@ -212,8 +268,7 @@ ${

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

213269 qout[base + (l + 1u - m)] = red[0].zw;

214270 }

215- }`

216- }

271+ }

217272 if (l + 2u > LMAX) { break; }

218273 let a0 = ab[base + (l + 2u - m)];

219274 var a1 = vec2f(0.0);

@@ -226,6 +281,7 @@ ${

226281 y1v[k] = a1.x * ctv[k] * t0 + a1.y * y1v[k];

227282 }

228283 l += 2u;

284+ }`

229285 }

230286 }

231287 `;