leg_analys: reduce with subgroupAdd, and retune wgAnalys for it - concept-collection/turing-sphere

concept-collection / turing-sphere

leg_analys: reduce with subgroupAdd, and retune wgAnalys for it

The analysis Legendre kernel reduced (c0, c1) across the workgroup with a shared-memory tree, once per l-pair: one barrier plus log2(wgAnalys) more, to emit two coefficients. The l-loop is serial and runs (lmax-m)/2 times, so those barriers sit on the critical path -- at lmax=127 the m=0 workgroup paid ~576 of them, while half the threads idled at every tree step. The file already noted the tree was a "portable stand-in for the CUDA warp shuffles" and the README listed subgroup reductions as future work; the `subgroups` feature makes it available. subgroupAdd reduces within a subgroup, and only the per-subgroup partials need combining, so it is two barriers per l-pair regardless of workgroup width. Both sit in uniform control flow. requestShtDevice now asks for the feature when the adapter offers it, and the tree is kept for when it does not. That inverts the wgAnalys trade-off. With the tree, a wider workgroup meant more barriers, and max(32, nlat/8) was the measured optimum. With subgroupAdd the barrier count no longer grows with width, so threads in flight -- (mmax+1) * wgAnalys, and analysis dispatches only mmax+1 workgroups -- becomes binding instead. 128 now wins at every grid. Both rules and their measurements are recorded at defaultWgAnalys. Measured on an RTX PRO 6000 Blackwell (driver 590.48.01, Dawn), bench:sht at lmax=63, 3 reps, ms per round trip: wg=256 + tree (as before) 0.090 - 0.097 + wgAnalys heuristic 0.067 + subgroupAdd 0.047 - 0.048 1.96x npm run bench --preset allencahn --lmax 127: 6090 -> 9960 steps/s, 1.64x. Round-trip accuracy is unchanged at 2.31e-6 relative L2, and test:node passes. A tuning() helper reads knobs from globalThis or the environment (SHT_WG_ANALYS, SHT_SUBGROUPS) so these can be A/B'd without editing code; that is how the numbers above were taken. Two things tried and rejected, recorded so they are not re-attempted blindly: parity folding measured 0.048 vs 0.047 here and 12% worse at 512x1024 -- the kernels are reduction-bound at these sizes, not arithmetic-bound, so halving the Legendre work moves nothing. And a barrier-free variant with one subgroup per workgroup was 1.52x faster than the tree at wgAnalys=32, but still lost to wg=128, because one subgroup per m is only 32 threads per m.

danfortunato <dan.fortunato@gmail.com> committed 20 days ago commit 53ff7c5857dc parent 4166b48 Browse files

2 changed files+104−7

src/sht/sht.tsmodified+67−2View file

@@ -35,7 +35,56 @@ export interface ShtOptions {

3535 }

3636

3737 const WG_SYNTH = 64;

38-const WG_ANALYS = 256;

38+

39+/**

40+ * Tuning knob, for A/B-ing a change without editing code. Reads globalThis

41+ * first (set it before creating a plan, as scripts/_ab.ts does), then the

42+ * environment, so `SHT_SUBGROUPS=0 npm run bench:sht` works too. `process` is

43+ * absent in the browser, where only the globalThis form applies.

44+ */

45+function tuning(name: string): unknown {

46+ const g = (globalThis as Record<string, unknown>)[name];

47+ if (g !== undefined) return g;

48+ const env = (globalThis as { process?: { env?: Record<string, string | undefined> } }).process?.env?.[name];

49+ if (env === undefined || env === '') return undefined;

50+ if (env === '1' || env === 'true') return true;

51+ if (env === '0' || env === 'false') return false;

52+ const n = Number(env);

53+ return Number.isFinite(n) ? n : env;

54+}

55+

56+/**

57+ * Workgroup size for the analysis Legendre reduction. The right answer differs

58+ * between the two reduction strategies, so it is chosen per strategy.

59+ *

60+ * Measured on an RTX PRO 6000 Blackwell (analysis, us). Shared-memory tree,

61+ * where each doubling of wgAnalys costs another barrier per l-pair:

62+ *

63+ * wgAnalys: 16 32 64 128 256

64+ * nlat=128 43.2 40.6 42.6 46.7 52.9 -> 32

65+ * nlat=256 104.0 85.6 87.8 89.9 100.0 -> 32

66+ * nlat=512 408.8 216.9 184.4 190.8 204.8 -> 64

67+ *

68+ * i.e. max(32, nlat/8). A flat 32 would be worse than the old default of 256 at

69+ * nlat=512, so it cannot be fitted on one grid. With subgroupAdd the barrier

70+ * count stops growing with wgAnalys and the picture inverts: threads in flight,

71+ * (mmax+1) * wgAnalys, becomes binding, since analysis dispatches only mmax+1

72+ * workgroups. 128 then wins at every grid (round trip, us):

73+ *

74+ * 128x256 37.8 (vs 38.3), 256x512 66.6 (vs 74.7), 512x1024 131.9 (vs 156.6)

75+ */

76+function defaultWgAnalys(nlat: number, limit: number, subgroups: boolean): number {

77+ if (subgroups) {

78+ // capped at nlat so small grids do not launch threads with no latitude to own

79+ let cap = 1;

80+ while (cap < nlat) cap *= 2;

81+ return Math.min(128, limit, cap);

82+ }

83+ const target = Math.max(32, nlat / 8);

84+ let wg = 1;

85+ while (wg < target) wg *= 2; // the tree reduction halves, so a power of two

86+ return Math.min(wg, limit);

87+}

3988

4089 async function makePipeline(

4190 device: GPUDevice,

@@ -170,7 +219,18 @@ export class ShtPlan {

170219 dev.queue.writeBuffer(this.bufTrig, 0, trig);

171220

172221 // --- shaders / pipelines ---

173- const legP = { lmax, mmax, nlat, wgSynth: WG_SYNTH, wgAnalys: WG_ANALYS };

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

223+ const wgAnalys =

224+ (tuning('SHT_WG_ANALYS') as number | undefined) ??

225+ defaultWgAnalys(nlat, dev.limits.maxComputeInvocationsPerWorkgroup, subgroups);

226+ const legP = {

227+ lmax,

228+ mmax,

229+ nlat,

230+ wgSynth: WG_SYNTH,

231+ wgAnalys,

232+ subgroups,

233+ };

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

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

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

@@ -359,7 +419,12 @@ export async function requestShtDevice(): Promise<GPUDevice> {

359419 if (!adapter) throw new Error('No WebGPU adapter available');

360420 // ask for a larger workgroup storage if the adapter offers it (bigger FFTs)

361421 const wgStorage = Math.min(adapter.limits.maxComputeWorkgroupStorageSize, 32768);

422+ // `subgroups` lets the analysis reduction use subgroupAdd instead of a

423+ // shared-memory tree (2 barriers per l-pair instead of 1 + log2(wgAnalys)).

424+ // Optional: ShtPlan falls back to the tree when it is not available.

425+ const features: GPUFeatureName[] = adapter.features.has('subgroups') ? ['subgroups'] : [];

362426 return adapter.requestDevice({

427+ requiredFeatures: features,

363428 requiredLimits: { maxComputeWorkgroupStorageSize: wgStorage },

364429 });

365430 }

src/sht/wgsl/leg.tsmodified+37−5View file

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

2323 nlat: number;

2424 wgSynth: number; // workgroup size for synthesis (threads over latitude)

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

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

27+ subgroups?: boolean;

2628 }

2729

2830 const BINDINGS = /* wgsl */ `

@@ -106,7 +108,14 @@ fn leg_synth(@builtin(global_invocation_id) gid: vec3u,

106108

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

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

109- return /* wgsl */ `

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' : ''}

110119 ${RESCALE_WGSL}

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

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

@@ -116,11 +125,15 @@ ${BINDINGS}

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

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

118127

119-var<workgroup> red: array<vec4f, ${p.wgAnalys}>;

128+var<workgroup> red: array<vec4f, ${redLen}>;

120129

121130 @compute @workgroup_size(${p.wgAnalys})

122131 fn leg_analys(@builtin(local_invocation_id) lid3: vec3u,

123- @builtin(workgroup_id) wid: 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+ }) {

124137 let lid = lid3.x;

125138 let m = wid.x;

126139 let base = m * (LMAX + 1u) - (m * (m - 1u)) / 2u;

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

167180 y1v[k] *= INV_SCALE;

168181 }

169182 }

170- // workgroup tree reduction of (c0, c1)

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)

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

172203 workgroupBarrier();

173204 var s = WG / 2u;

@@ -181,7 +212,8 @@ fn leg_analys(@builtin(local_invocation_id) lid3: vec3u,

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

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

183214 }

184- }

215+ }`

216+ }

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

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

187219 var a1 = vec2f(0.0);